from __future__ import annotations

import pandas as pd
from rdkit import Chem
from rdkit.Chem import Crippen, Descriptors, Lipinski, QED

from ._core import classify_macrolactone, ensure_mol, find_macrolactone_candidates
from .models import MacrocycleClassificationResult


class MacroLactoneAnalyzer:
    """Identify valid 12-20 membered macrolactones."""

    def get_valid_ring_sizes(self, mol_input: str | Chem.Mol) -> list[int]:
        mol, _ = ensure_mol(mol_input)
        candidates = find_macrolactone_candidates(mol)
        return sorted({candidate.ring_size for candidate in candidates})

    def classify_macrocycle(
        self,
        mol_input: str | Chem.Mol,
        ring_size: int | None = None,
    ) -> MacrocycleClassificationResult:
        mol, smiles = ensure_mol(mol_input)
        return classify_macrolactone(mol, smiles=smiles, ring_size=ring_size)

    def analyze_molecule(
        self,
        mol_input: str | Chem.Mol,
        ring_size: int | None = None,
    ) -> dict:
        return self.classify_macrocycle(mol_input, ring_size=ring_size).to_dict()

    def analyze_many(
        self,
        smiles_list: list[str],
        ring_range: range = range(12, 21),
    ) -> dict:
        classification_counts = {
            "standard_macrolactone": 0,
            "non_standard_macrocycle": 0,
            "not_macrolactone": 0,
        }
        ring_size_counts = {ring_size: 0 for ring_size in ring_range}
        results: list[dict] = []

        for smiles in smiles_list:
            classification = self.classify_macrocycle(smiles)
            classification_counts[classification.classification] += 1
            if (
                classification.classification == "standard_macrolactone"
                and classification.ring_size in ring_size_counts
            ):
                ring_size_counts[classification.ring_size] += 1
            results.append(classification.to_dict())

        return {
            "total": len(smiles_list),
            "classification_counts": classification_counts,
            "ring_size_counts": ring_size_counts,
            "results": results,
        }

    def classify_dataframe(
        self,
        dataframe: pd.DataFrame,
        smiles_column: str = "smiles",
        id_column: str | None = None,
        ring_range: range = range(12, 21),
    ) -> tuple[dict[int, pd.DataFrame], pd.DataFrame]:
        grouped_rows: dict[int, list[dict]] = {ring_size: [] for ring_size in ring_range}
        rejected_rows: list[dict] = []

        for index, row in dataframe.iterrows():
            base_row = row.to_dict()
            if id_column is None and "id" not in base_row:
                base_row["id"] = f"row_{index}"

            classification = self.classify_macrocycle(base_row[smiles_column])
            enriched_row = {
                **base_row,
                **classification.to_dict(),
            }
            if (
                classification.classification == "standard_macrolactone"
                and classification.ring_size in grouped_rows
            ):
                grouped_rows[classification.ring_size].append(enriched_row)
            else:
                rejected_rows.append(enriched_row)

        grouped_frames = {
            ring_size: pd.DataFrame(rows)
            for ring_size, rows in grouped_rows.items()
            if rows
        }
        return grouped_frames, pd.DataFrame(rejected_rows)

    def match_dynamic_smarts(self, smiles: str, ring_size: int) -> list[int] | None:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        query = Chem.MolFromSmarts(f"[r{ring_size}]([#8][#6](=[#8]))")
        if query is None:
            return None
        matches = mol.GetSubstructMatches(query)
        return list(matches[0]) if matches else None

    def calculate_properties(self, smiles: str) -> dict[str, float] | None:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        return {
            "molecular_weight": Descriptors.MolWt(mol),
            "logp": Crippen.MolLogP(mol),
            "qed": QED.qed(mol),
            "tpsa": Descriptors.TPSA(mol),
            "num_atoms": float(mol.GetNumAtoms()),
            "num_heavy_atoms": float(mol.GetNumHeavyAtoms()),
            "num_h_donors": float(Lipinski.NumHDonors(mol)),
            "num_h_acceptors": float(Lipinski.NumHAcceptors(mol)),
            "num_rotatable_bonds": float(Lipinski.NumRotatableBonds(mol)),
        }
