# Copyright 2024-2025 SpectrixRD
# Licensed under the Apache License, Version 2.0
# See LICENSE file in the project root for full license text

"""
QPHAROS Data Models
"""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class DockingResult(BaseModel):
    """Result from QPHAROS molecular docking"""

    ligand_smiles: str = Field(..., description="SMILES string of ligand")
    receptor_pdb: str = Field(..., description="PDB ID of receptor")

    # Core docking results
    binding_affinity: Optional[float] = Field(None, description="Binding affinity in kcal/mol")
    ki: Optional[float] = Field(None, description="Inhibition constant Ki in nM")
    ic50: Optional[float] = Field(None, description="IC50 in nM")
    docking_score: Optional[float] = Field(None, description="Docking score")

    # Molecular interactions
    h_bonds: Optional[int] = Field(None, description="Number of hydrogen bonds")
    hydrophobic_contacts: Optional[int] = Field(None, description="Number of hydrophobic contacts")
    salt_bridges: Optional[int] = Field(None, description="Number of salt bridges")
    pi_stacking: Optional[int] = Field(None, description="Number of π-π stacking interactions")

    # Drug-likeness
    lipinski_pass: Optional[bool] = Field(None, description="Lipinski Rule of 5 compliance")
    qed_score: Optional[float] = Field(None, description="QED drug-likeness score (0-1)")
    sa_score: Optional[float] = Field(None, description="Synthetic accessibility score (1-10)")
    pains_alerts: Optional[int] = Field(None, description="Number of PAINS alerts")

    # Quantum metadata
    backend: Optional[str] = Field(None, description="Quantum backend used")
    shots: Optional[int] = Field(None, description="Number of quantum shots")
    job_id: Optional[str] = Field(None, description="IBM Quantum job ID")
    qec_enabled: bool = Field(False, description="Whether QEC was used")
    quantum_counts: Optional[Dict[str, int]] = Field(None, description="Quantum measurement counts")

    success: bool = Field(..., description="Whether execution succeeded")
    error_message: Optional[str] = Field(None, description="Error message if failed")

    @classmethod
    def from_bioql_result(cls, result: Any) -> "DockingResult":
        """Create DockingResult from BioQL quantum result"""

        bio = getattr(result, "bio_interpretation", {}) or {}
        metadata = getattr(result, "metadata", {}) or {}

        return cls(
            ligand_smiles=metadata.get("ligand_smiles", "unknown"),
            receptor_pdb=metadata.get("receptor_pdb", "unknown"),
            binding_affinity=bio.get("binding_affinity"),
            ki=bio.get("ki"),
            ic50=bio.get("ic50"),
            docking_score=bio.get("docking_score"),
            h_bonds=bio.get("h_bonds"),
            hydrophobic_contacts=bio.get("hydrophobic_contacts"),
            salt_bridges=bio.get("salt_bridges"),
            pi_stacking=bio.get("pi_stacking"),
            lipinski_pass=bio.get("lipinski_pass"),
            qed_score=bio.get("qed_score"),
            sa_score=bio.get("sa_score"),
            pains_alerts=bio.get("pains_alerts"),
            backend=getattr(result, "backend_name", None),
            shots=metadata.get("shots"),
            job_id=getattr(result, "job_id", None),
            qec_enabled=metadata.get("qec_enabled", False),
            quantum_counts=getattr(result, "counts", None),
            success=getattr(result, "success", False),
            error_message=getattr(result, "error_message", None),
        )

    def __str__(self) -> str:
        lines = [
            "QPHAROS Docking Result",
            "=" * 60,
            f"Ligand: {self.ligand_smiles}",
            f"Receptor: {self.receptor_pdb}",
            "",
        ]

        if self.binding_affinity:
            lines.append(f"Binding Affinity: {self.binding_affinity:.2f} kcal/mol")

        if self.ki:
            lines.append(f"Ki: {self.ki:.2f} nM")

        if self.ic50:
            lines.append(f"IC50: {self.ic50:.2f} nM")

        if self.h_bonds is not None:
            lines.append(f"\\nInteractions:")
            lines.append(f"  H-bonds: {self.h_bonds}")
            lines.append(f"  Hydrophobic: {self.hydrophobic_contacts or 0}")

        if self.qed_score:
            lines.append(f"\\nDrug-likeness:")
            lines.append(f"  QED Score: {self.qed_score:.3f}")
            lines.append(f"  Lipinski Pass: {'✓' if self.lipinski_pass else '✗'}")

        if self.job_id:
            lines.append(f"\\nQuantum Job ID: {self.job_id}")

        return "\\n".join(lines)


class GeneratedMolecule(BaseModel):
    """A molecule generated by QPHAROS QGAN"""

    smiles: str = Field(..., description="SMILES string")
    score: float = Field(..., description="Generation score (0-1)")
    binding_affinity: Optional[float] = Field(None, description="Predicted binding affinity")
    qed: Optional[float] = Field(None, description="QED drug-likeness")
    sa_score: Optional[float] = Field(None, description="Synthetic accessibility")
    lipinski_pass: Optional[bool] = Field(None, description="Lipinski compliance")


class DrugDesignResult(BaseModel):
    """Result from QPHAROS drug design (QGAN)"""

    target_protein: str = Field(..., description="Target protein PDB ID")
    molecules: List[GeneratedMolecule] = Field(
        default_factory=list, description="Generated molecules"
    )

    backend: Optional[str] = Field(None, description="Quantum backend")
    shots: Optional[int] = Field(None, description="Quantum shots")
    job_id: Optional[str] = Field(None, description="IBM Quantum job ID")

    success: bool = Field(..., description="Execution success")
    error_message: Optional[str] = Field(None, description="Error if failed")

    @classmethod
    def from_bioql_result(cls, result: Any) -> "DrugDesignResult":
        """Create DrugDesignResult from BioQL result"""

        bio = getattr(result, "bio_interpretation", {}) or {}
        metadata = getattr(result, "metadata", {}) or {}

        # Parse generated molecules from bio_interpretation
        molecules = []
        if "generated_molecules" in bio:
            for mol_data in bio["generated_molecules"]:
                molecules.append(GeneratedMolecule(**mol_data))

        return cls(
            target_protein=metadata.get("target_protein", "unknown"),
            molecules=molecules,
            backend=getattr(result, "backend_name", None),
            shots=metadata.get("shots"),
            job_id=getattr(result, "job_id", None),
            success=getattr(result, "success", False),
            error_message=getattr(result, "error_message", None),
        )


class ADMETResult(BaseModel):
    """Result from QPHAROS ADMET prediction"""

    smiles: str = Field(..., description="SMILES string")

    # Absorption
    caco2_permeability: Optional[float] = Field(None, description="Caco-2 permeability")
    hia: Optional[float] = Field(None, description="Human Intestinal Absorption (%)")

    # Distribution
    vdss: Optional[float] = Field(None, description="Volume of distribution (L/kg)")
    ppb: Optional[float] = Field(None, description="Plasma protein binding (%)")

    # Metabolism
    cyp3a4_substrate: Optional[bool] = Field(None, description="CYP3A4 substrate")
    cyp2d6_substrate: Optional[bool] = Field(None, description="CYP2D6 substrate")
    cyp2c9_substrate: Optional[bool] = Field(None, description="CYP2C9 substrate")

    # Excretion
    clearance: Optional[float] = Field(None, description="Clearance (mL/min/kg)")
    half_life: Optional[float] = Field(None, description="Half-life (hours)")

    # Toxicity
    herg_inhibition: Optional[bool] = Field(None, description="hERG inhibition risk")
    ames_mutagenicity: Optional[bool] = Field(None, description="AMES mutagenicity")
    ld50: Optional[float] = Field(None, description="LD50 (mg/kg)")

    # Drug-likeness
    lipinski_pass: Optional[bool] = Field(None, description="Lipinski Rule of 5")
    qed_score: Optional[float] = Field(None, description="QED score")
    sa_score: Optional[float] = Field(None, description="Synthetic accessibility")
    pains_alerts: Optional[int] = Field(None, description="PAINS alerts")

    # Quantum metadata
    backend: Optional[str] = Field(None, description="Quantum backend")
    shots: Optional[int] = Field(None, description="Quantum shots")
    job_id: Optional[str] = Field(None, description="Job ID")

    success: bool = Field(..., description="Success flag")
    error_message: Optional[str] = Field(None, description="Error message")

    @classmethod
    def from_bioql_result(cls, result: Any) -> "ADMETResult":
        """Create ADMETResult from BioQL result"""

        bio = getattr(result, "bio_interpretation", {}) or {}
        metadata = getattr(result, "metadata", {}) or {}

        return cls(
            smiles=metadata.get("smiles", "unknown"),
            caco2_permeability=bio.get("caco2_permeability"),
            hia=bio.get("hia"),
            vdss=bio.get("vdss"),
            ppb=bio.get("ppb"),
            cyp3a4_substrate=bio.get("cyp3a4_substrate"),
            cyp2d6_substrate=bio.get("cyp2d6_substrate"),
            cyp2c9_substrate=bio.get("cyp2c9_substrate"),
            clearance=bio.get("clearance"),
            half_life=bio.get("half_life"),
            herg_inhibition=bio.get("herg_inhibition"),
            ames_mutagenicity=bio.get("ames_mutagenicity"),
            ld50=bio.get("ld50"),
            lipinski_pass=bio.get("lipinski_pass"),
            qed_score=bio.get("qed_score"),
            sa_score=bio.get("sa_score"),
            pains_alerts=bio.get("pains_alerts"),
            backend=getattr(result, "backend_name", None),
            shots=metadata.get("shots"),
            job_id=getattr(result, "job_id", None),
            success=getattr(result, "success", False),
            error_message=getattr(result, "error_message", None),
        )
