from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import List, Optional

class MoleculeInput(BaseModel):
    """Molecular input for analysis."""
    smiles: str = Field(description="SMILES notation of molecule")
    name: Optional[str] = Field(description="Common name if known")

@tool
def analyze_molecule(smiles: str) -> str:
    """Analyze molecular properties using RDKit.
    
    Args:
        smiles: SMILES notation of the molecule
    """
    try:
        from rdkit import Chem
        from rdkit.Chem import Descriptors, Lipinski, Draw
        
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return f"Invalid SMILES: {smiles}"
        
        # Calculate properties
        mw = Descriptors.MolWt(mol)
        logp = Descriptors.MolLogP(mol)
        hbd = Lipinski.NumHDonors(mol)
        hba = Lipinski.NumHAcceptors(mol)
        tpsa = Descriptors.TPSA(mol)
        rotatable = Lipinski.NumRotatableBonds(mol)
        
        # Lipinski check
        lipinski_pass = mw <= 500 and logp <= 5 and hbd <= 5 and hba <= 10
        
        return f"""
Molecular Analysis:
- Name: {Chem.MolToSmiles(mol)}
- Molecular Weight: {mw:.1f} g/mol
- LogP: {logp:.2f}
- H-Bond Donors: {hbd}
- H-Bond Acceptors: {hba}
- TPSA: {tpsa:.1f} Ų
- Rotatable Bonds: {rotatable}
- Lipinski Compliant: {'Yes ✓' if lipinski_pass else 'No ✗'}
"""
    except ImportError:
        return "RDKit not installed. Install with: pip install rdkit-pypi"
    except Exception as e:
        return f"Analysis error: {str(e)}"


@tool
def search_similar_compounds(smiles: str, top_k: int = 5) -> str:
    """Search for similar compounds using Tanimoto similarity.
    
    Args:
        smiles: Reference molecule SMILES
        top_k: Number of similar compounds to return
    """
    try:
        from rdkit import Chem
        from rdkit.Chem import AllChem, DataStructs
        import numpy as np
        
        ref_mol = Chem.MolFromSmiles(smiles)
        if ref_mol is None:
            return f"Invalid SMILES: {smiles}"
        
        # Generate fingerprint for reference
        ref_fp = AllChem.GetMorganFingerprintAsBitVect(ref_mol, 2, nBits=1024)
        
        # Mock database (in production, use real database)
        mock_db = [
            {"smiles": "CCO", "name": "Ethanol", "activity": "low"},
            {"smiles": "CC(=O)O", "name": "Acetic acid", "activity": "medium"},
            {"smiles": "c1ccc(CC(=O)O)cc1", "name": "Phenylacetic acid", "activity": "high"},
        ]
        
        results = []
        for compound in mock_db:
            mol = Chem.MolFromSmiles(compound["smiles"])
            if mol:
                fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024)
                similarity = DataStructs.TanimotoSimilarity(ref_fp, fp)
                results.append({**compound, "similarity": similarity})
        
        # Sort by similarity
        results.sort(key=lambda x: x["similarity"], reverse=True)
        
        output = f"Similar compounds to {smiles}:\n\n"
        for i, r in enumerate(results[:top_k], 1):
            output += f"{i}. {r['name']} (Sim: {r['similarity']:.3f})\n"
            output += f"   SMILES: {r['smiles']}\n"
            output += f"   Activity: {r['activity']}\n\n"
        
        return output
    except ImportError:
        return "RDKit not installed"
    except Exception as e:
        return f"Search error: {str(e)}"


@tool
def predict_admet(smiles: str) -> str:
    """Predict ADMET properties for a molecule.
    
    Args:
        smiles: SMILES notation of molecule
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Predict ADMET properties for this molecule:
    
    SMILES: {smiles}
    
    Provide predictions for:
    1. Absorption (oral bioavailability, Caco-2 permeability)
    2. Distribution (Plasma protein binding, Vd)
    3. Metabolism (CYP inhibition, half-life)
    4. Excretion (renal clearance)
    5. Toxicity (hERG, hepatotoxicity, mutagenicity)
    
    Format as structured report:"""
    
    response = llm.invoke(prompt)
    return response.content


@tool
def suggest_modifications(smiles: str, goal: str) -> str:
    """Suggest molecular modifications to improve properties.
    
    Args:
        smiles: Current molecule SMILES
        goal: Optimization goal (e.g., "improve solubility", "reduce toxicity")
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Suggest molecular modifications for:
    
    Molecule: {smiles}
    Goal: {goal}
    
    Provide:
    1. Specific structural modifications
    2. Expected impact on properties
    3. Synthetic feasibility
    4. Alternative scaffolds
    
    Use medicinal chemistry principles:"""
    
    response = llm.invoke(prompt)
    return response.content

from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

class DrugDiscoveryState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    query_smiles: str
    analysis: str
    similar_compounds: str
    admet: str
    optimization: str
    recommendations: str

def analyze_node(state: DrugDiscoveryState):
    """Analyze molecular properties."""
    result = analyze_molecule.invoke({"smiles": state["query_smiles"]})
    return {"analysis": result}

def search_node(state: DrugDiscoveryState):
    """Search for similar compounds."""
    result = search_similar_compounds.invoke({
        "smiles": state["query_smiles"],
        "top_k": 5,
    })
    return {"similar_compounds": result}

def admet_node(state: DrugDiscoveryState):
    """Predict ADMET properties."""
    result = predict_admet.invoke({"smiles": state["query_smiles"]})
    return {"admet": result}

def optimize_node(state: DrugDiscoveryState):
    """Suggest lead optimizations."""
    result = suggest_modifications.invoke({
        "smiles": state["query_smiles"],
        "goal": "improve drug-likeness and reduce toxicity",
    })
    return {"optimization": result}

def recommend_node(state: DrugDiscoveryState):
    """Generate final recommendations."""
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    prompt = f"""Summarize drug discovery analysis:
    
    Molecule: {state['query_smiles']}
    
    Properties:
    {state['analysis']}
    
    Similar Compounds:
    {state['similar_compounds']}
    
    ADMET:
    {state['admet']}
    
    Optimization Suggestions:
    {state['optimization']}
    
    Provide final recommendations for this compound:"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    return {"recommendations": response.content, "messages": [response]}

# Build graph
workflow = StateGraph(DrugDiscoveryState)
workflow.add_node("analyze", analyze_node)
workflow.add_node("search", search_node)
workflow.add_node("admet", admet_node)
workflow.add_node("optimize", optimize_node)
workflow.add_node("recommend", recommend_node)

workflow.set_entry_point("analyze")
workflow.add_edge("analyze", "search")
workflow.add_edge("search", "admet")
workflow.add_edge("admet", "optimize")
workflow.add_edge("optimize", "recommend")
workflow.add_edge("recommend", END)

app = workflow.compile()

# Run
result = app.invoke({
    "messages": [HumanMessage(content="Analyze aspirin for drug discovery")],
    "query_smiles": "CC(=O)Oc1ccccc1C(=O)O",
    "analysis": "",
    "similar_compounds": "",
    "admet": "",
    "optimization": "",
    "recommendations": "",
})
print(result["recommendations"])