src/chemgraph/tools/cheminformatics_core.py

"""Pure-Python cheminformatics helpers (no LangChain / MCP decorators).

Provides a single implementation for PubChem lookups and RDKit
SMILES-to-3D conversion, used by both the LangChain ``@tool`` wrappers
in :mod:`cheminformatics_tools` and the MCP wrappers in
:mod:`chemgraph.mcp.mcp_tools`.
"""

from __future__ import annotations

import os
import re
from typing import Literal

import pubchempy as pcp

from chemgraph.schemas.atomsdata import AtomsData
from chemgraph.tools.ase_core import _resolve_path


# ---------------------------------------------------------------------------
# SMILES → 3D coordinates (single implementation)
# ---------------------------------------------------------------------------


def smiles_to_3d(
    smiles: str, seed: int = 2025
) -> tuple[list[int], list[list[float]]]:
    """Convert a SMILES string to 3D coordinates via RDKit.

    Parameters
    ----------
    smiles : str
        SMILES string representation of the molecule.
    seed : int, optional
        Random seed for reproducible 3D embedding, by default 2025.

    Returns
    -------
    tuple[list[int], list[list[float]]]
        ``(atomic_numbers, positions)`` where *positions* is a list of
        ``[x, y, z]`` lists in Angstroms.

    Raises
    ------
    ValueError
        If the SMILES string is invalid or 3D generation/optimization fails.
    """
    from rdkit import Chem
    from rdkit.Chem import AllChem

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError("Invalid SMILES string.")

    mol = Chem.AddHs(mol)
    if AllChem.EmbedMolecule(mol, randomSeed=seed) != 0:
        raise ValueError("Failed to generate 3D coordinates.")
    if AllChem.UFFOptimizeMolecule(mol) != 0:
        raise ValueError("Failed to optimize 3D geometry.")

    conf = mol.GetConformer()
    numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
    positions = [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())]
    return numbers, positions


# ---------------------------------------------------------------------------
# PubChem name → SMILES
# ---------------------------------------------------------------------------

_MIXTURE_LIKE_NAMES: dict[str, str] = {
    "vinegar": (
        "Vinegar is usually an aqueous acetic-acid solution, not a single "
        "pure molecule."
    ),
    "bleach": (
        "Bleach is usually an aqueous hypochlorite product, not a single pure "
        "molecule."
    ),
    "rubbing alcohol": (
        "Rubbing alcohol is a solution/product name, not a unique pure molecule."
    ),
    "battery acid": (
        "Battery acid is usually aqueous sulfuric acid, not a single pure molecule."
    ),
    "lye": "Lye can refer to sodium hydroxide or potassium hydroxide solutions.",
    "peroxide": (
        "Peroxide is often used as a product nickname and may require "
        "concentration/component clarification."
    ),
}

_MIXTURE_HINT_TOKENS = {
    "solution",
    "mixture",
    "household",
    "commercial",
    "extract",
    "solvent",
    "cleaner",
    "grade",
}

_HYDRATE_SOLVATE_TOKENS = {
    "hydrate",
    "monohydrate",
    "dihydrate",
    "trihydrate",
    "tetrahydrate",
    "pentahydrate",
    "hexahydrate",
    "solvate",
    "hemihydrate",
}

_SALT_ADDUCT_TOKENS = {
    "salt",
    "hydrochloride",
    "sodium",
    "potassium",
    "lithium",
    "calcium",
    "magnesium",
    "ammonium",
    "chloride",
    "bromide",
    "iodide",
    "acetate",
    "sulfate",
    "phosphate",
}

_STEREO_TOKENS = {
    "stereo",
    "stereoisomer",
    "enantiomer",
    "chiral",
    "cis",
    "trans",
    "r",
    "s",
    "e",
    "z",
}


def resolve_molecule_identity_core(name: str, max_candidates: int = 5) -> dict:
    """Resolve a molecule name to a PubChem-backed identity record.

    This richer resolver preserves the legacy SMILES behavior while exposing
    provenance, candidate summaries, and a conservative credibility score. The
    agent can then decide whether to proceed, ask for clarification, or state an
    explicit representative-molecule assumption.
    """
    if not name or not str(name).strip():
        raise ValueError("Parameter 'name' must be a non-empty string.")

    input_name = str(name).strip()
    comps = pcp.get_compounds(input_name, "name")
    if not comps:
        raise ValueError(f"No PubChem compound found for name: {name!r}")

    query_flags = _query_identity_flags(input_name)
    mixture_note = _mixture_note(input_name)
    candidates = [
        _compound_candidate(
            input_name,
            compound,
            index,
            len(comps),
            bool(mixture_note),
            query_flags,
        )
        for index, compound in enumerate(comps[:max_candidates])
    ]
    candidates = [candidate for candidate in candidates if candidate.get("smiles")]
    if not candidates:
        raise ValueError(f"PubChem returned an empty SMILES for {name!r}.")

    selected = candidates[0]
    identity_flags = dict(selected.get("identity_flags") or {})
    identity_flags.update(
        {
            "mixture_like_name": bool(mixture_note),
            "ambiguous": _is_ambiguous(candidates, bool(mixture_note)),
        }
    )
    warnings = _identity_warnings(input_name, selected, identity_flags, mixture_note)
    requires_clarification = bool(warnings) or selected["credibility_score"] < 0.5
    return {
        "status": "needs_clarification" if requires_clarification else "resolved",
        "input_name": input_name,
        "resolved_name": selected.get("iupac_name") or input_name,
        "smiles": selected["smiles"],
        "canonical_smiles": selected.get("canonical_smiles"),
        "isomeric_smiles": selected.get("isomeric_smiles"),
        "connectivity_smiles": selected.get("connectivity_smiles"),
        "molecular_formula": selected.get("molecular_formula"),
        "inchikey": selected.get("inchikey"),
        "source": "PubChem",
        "cid": selected.get("cid"),
        "candidates": candidates,
        "selected_candidate_index": 0,
        "confidence_score": selected["confidence_score"],
        "credibility_score": selected["credibility_score"],
        "score_breakdown": selected.get("score_breakdown", {}),
        "identity_flags": identity_flags,
        "resolver_provenance": {
            "resolver": "PubChemPy",
            "namespace": "name",
            "candidate_count": len(comps),
            "returned_candidate_count": len(candidates),
            "max_candidates": max_candidates,
            "selection": "top-ranked PubChem candidate",
        },
        "ambiguity_flag": identity_flags["ambiguous"],
        "mixture_flag": identity_flags["mixture_like_name"],
        "is_mixture": identity_flags["mixture_like_name"],
        "requires_clarification": requires_clarification,
        "needs_clarification": requires_clarification,
        "representative_of": input_name if requires_clarification else None,
        "selection_reason": _selection_reason(
            input_name, selected, candidates, mixture_note
        ),
        "warnings": warnings,
        "warning": "; ".join(warnings) if warnings else "",
    }


def molecule_name_to_smiles_core(name: str) -> str:
    """Resolve a molecule name to its canonical SMILES via PubChem.

    Parameters
    ----------
    name : str
        Common or IUPAC molecule name.

    Returns
    -------
    str
        Canonical SMILES string.

    Raises
    ------
    ValueError
        If no PubChem match is found or the returned SMILES is empty.
    """
    return resolve_molecule_identity_core(name)["smiles"]


def _compound_candidate(
    input_name: str,
    compound: object,
    index: int,
    total_count: int,
    mixture_like: bool,
    query_flags: dict,
) -> dict:
    isomeric_smiles = getattr(compound, "isomeric_smiles", None)
    canonical_smiles = getattr(compound, "canonical_smiles", None)
    connectivity_smiles = getattr(compound, "connectivity_smiles", None)
    smiles = isomeric_smiles or connectivity_smiles or canonical_smiles
    iupac_name = getattr(compound, "iupac_name", None)
    cid = getattr(compound, "cid", None)
    formula = getattr(compound, "molecular_formula", None)
    inchikey = getattr(compound, "inchikey", None)
    synonyms = _safe_synonyms(compound)
    candidate_flags = _candidate_identity_flags(
        smiles=smiles,
        isomeric_smiles=isomeric_smiles,
        canonical_smiles=canonical_smiles,
        input_flags=query_flags,
    )
    score = _candidate_score(
        input_name=input_name,
        iupac_name=iupac_name,
        synonyms=synonyms,
        index=index,
        total_count=total_count,
        mixture_like=mixture_like,
        candidate_flags=candidate_flags,
    )
    return {
        "rank": index + 1,
        "cid": cid,
        "canonical_smiles": smiles,
        "smiles": smiles,
        "isomeric_smiles": isomeric_smiles,
        "connectivity_smiles": connectivity_smiles,
        "pubchem_canonical_smiles": canonical_smiles,
        "molecular_formula": formula,
        "inchikey": inchikey,
        "iupac_name": iupac_name,
        "synonyms_sample": synonyms[:6],
        "confidence_score": score,
        "credibility_score": score,
        "identity_flags": candidate_flags,
        "candidate_flags": candidate_flags,
        "score_breakdown": _score_breakdown(
            input_name=input_name,
            iupac_name=iupac_name,
            synonyms=synonyms,
            index=index,
            total_count=total_count,
            mixture_like=mixture_like,
            candidate_flags=candidate_flags,
            score=score,
        ),
        "source": "PubChem",
    }


def _safe_synonyms(compound: object) -> list[str]:
    try:
        synonyms = getattr(compound, "synonyms", None) or []
    except Exception:
        synonyms = []
    return [str(value) for value in synonyms if value][:12]


def _candidate_score(
    *,
    input_name: str,
    iupac_name: str | None,
    synonyms: list[str],
    index: int,
    total_count: int,
    mixture_like: bool,
    candidate_flags: dict,
) -> float:
    if mixture_like:
        return 0.35

    normalized_input = _normalize_identity_text(input_name)
    labels = [iupac_name or "", *synonyms]
    normalized_labels = {_normalize_identity_text(label) for label in labels if label}
    if normalized_input in normalized_labels:
        base = 0.95
    elif any(
        normalized_input and normalized_input in label for label in normalized_labels
    ):
        base = 0.85
    elif total_count == 1:
        base = 0.8
    else:
        base = 0.68

    rank_penalty = min(index * 0.08, 0.3)
    structural_penalty = 0.0
    if candidate_flags.get("multi_fragment_smiles"):
        structural_penalty += 0.18
    if candidate_flags.get("stereo_requested") and not candidate_flags.get(
        "stereo_preserved"
    ):
        structural_penalty += 0.18
    if candidate_flags.get("hydrate_or_solvate_name"):
        structural_penalty += 0.08
    if candidate_flags.get("salt_or_adduct_name"):
        structural_penalty += 0.08
    return round(max(0.1, min(0.99, base - rank_penalty - structural_penalty)), 2)


def _normalize_identity_text(text: str | None) -> str:
    return re.sub(r"[^a-z0-9]+", " ", str(text or "").lower()).strip()


def _mixture_note(name: str) -> str:
    normalized = _normalize_identity_text(name)
    if normalized in _MIXTURE_LIKE_NAMES:
        return _MIXTURE_LIKE_NAMES[normalized]
    if any(token in normalized.split() for token in _MIXTURE_HINT_TOKENS):
        return (
            "The name appears to describe a product, mixture, or solution rather "
            "than a single pure molecule."
        )
    return ""


def _query_identity_flags(name: str) -> dict:
    tokens = set(_normalize_identity_text(name).split())
    return {
        "hydrate_or_solvate_name": bool(tokens & _HYDRATE_SOLVATE_TOKENS),
        "salt_or_adduct_name": bool(tokens & _SALT_ADDUCT_TOKENS),
        "stereo_requested": bool(tokens & _STEREO_TOKENS),
    }


def _candidate_identity_flags(
    *,
    smiles: str | None,
    isomeric_smiles: str | None,
    canonical_smiles: str | None,
    input_flags: dict,
) -> dict:
    smiles_text = smiles or ""
    isomeric_text = isomeric_smiles or ""
    stereo_preserved = bool(
        isomeric_text and any(marker in isomeric_text for marker in ("@", "/", "\\"))
    )
    return {
        **input_flags,
        "multi_fragment_smiles": "." in smiles_text,
        "stereo_preserved": stereo_preserved,
        "isomeric_smiles_available": bool(isomeric_smiles),
        "canonical_smiles_available": bool(canonical_smiles),
    }


def _identity_warnings(
    input_name: str,
    selected: dict,
    identity_flags: dict,
    mixture_note: str,
) -> list[str]:
    warnings: list[str] = []
    if mixture_note:
        warnings.append(mixture_note)
    if identity_flags.get("multi_fragment_smiles"):
        warnings.append(
            "The selected PubChem structure has multiple disconnected fragments; "
            "clarify whether to model the full salt/adduct or a neutral component."
        )
    if identity_flags.get("hydrate_or_solvate_name"):
        warnings.append(
            "The name suggests a hydrate or solvate; clarify whether waters/solvent "
            "should be included in the calculation."
        )
    if identity_flags.get("salt_or_adduct_name"):
        warnings.append(
            "The name suggests a salt or adduct; clarify the modeled component, "
            "charge state, and counterion handling."
        )
    if identity_flags.get("stereo_requested") and not identity_flags.get(
        "stereo_preserved"
    ):
        warnings.append(
            f"The query {input_name!r} appears stereochemistry-sensitive, but the "
            "selected SMILES does not preserve explicit stereochemistry."
        )
    if float(selected.get("credibility_score", 0.0)) < 0.5:
        warnings.append(
            "The top PubChem candidate has low credibility for this input name; "
            "confirm the identity before calculation."
        )
    return warnings


def _score_breakdown(
    *,
    input_name: str,
    iupac_name: str | None,
    synonyms: list[str],
    index: int,
    total_count: int,
    mixture_like: bool,
    candidate_flags: dict,
    score: float,
) -> dict:
    normalized_input = _normalize_identity_text(input_name)
    normalized_labels = [
        _normalize_identity_text(label) for label in [iupac_name or "", *synonyms] if label
    ]
    return {
        "final_score": score,
        "rank": index + 1,
        "total_pubchem_matches": total_count,
        "exact_label_match": normalized_input in normalized_labels,
        "substring_label_match": any(
            normalized_input and normalized_input in label
            for label in normalized_labels
        ),
        "mixture_like_penalty": mixture_like,
        "multi_fragment_penalty": bool(candidate_flags.get("multi_fragment_smiles")),
        "salt_or_adduct_penalty": bool(candidate_flags.get("salt_or_adduct_name")),
        "hydrate_or_solvate_penalty": bool(
            candidate_flags.get("hydrate_or_solvate_name")
        ),
        "stereo_missing_penalty": bool(
            candidate_flags.get("stereo_requested")
            and not candidate_flags.get("stereo_preserved")
        ),
    }


def _is_ambiguous(candidates: list[dict], requires_clarification: bool) -> bool:
    if requires_clarification:
        return True
    if len(candidates) <= 1:
        return False
    top_score = float(candidates[0].get("credibility_score", 0.0))
    return top_score < 0.9


def _selection_reason(
    input_name: str,
    selected: dict,
    candidates: list[dict],
    mixture_note: str,
) -> str:
    if mixture_note:
        return (
            f"PubChem returned a candidate for {input_name!r}, but the input looks "
            "like a mixture or product name. Clarify the component/composition or "
            "state an explicit representative molecule before calculation."
        )
    if len(candidates) == 1:
        return "Single PubChem candidate was returned for the input name."
    return (
        "Top PubChem candidate selected. Review candidates and credibility_score "
        "when the name is ambiguous."
    )


# ---------------------------------------------------------------------------
# SMILES → coordinate file
# ---------------------------------------------------------------------------


def smiles_to_coordinate_file_core(
    smiles: str,
    output_file: str = "molecule.xyz",
    seed: int = 2025,
    fmt: Literal["xyz"] = "xyz",
) -> dict:
    """Convert a SMILES string to a coordinate file on disk.

    Parameters
    ----------
    smiles : str
        SMILES string representation of the molecule.
    output_file : str, optional
        Path to save the output coordinate file.
    seed : int, optional
        Random seed for RDKit 3D structure generation, by default 2025.
    fmt : {"xyz"}, optional
        Output format.  Only ``"xyz"`` is supported currently.

    Returns
    -------
    dict
        ``{"ok": True, "artifact": "coordinate_file", "path": ...,
        "smiles": ..., "natoms": ...}``

    Raises
    ------
    ValueError
        If the SMILES string is invalid or 3D generation fails.
    """
    from ase import Atoms
    from ase.io import write as ase_write

    numbers, positions = smiles_to_3d(smiles, seed=seed)
    atoms = Atoms(numbers=numbers, positions=positions)

    final_output_file = _resolve_path(output_file)
    ase_write(final_output_file, atoms)

    return {
        "ok": True,
        "artifact": "coordinate_file",
        "path": os.path.abspath(final_output_file),
        "smiles": smiles,
        "natoms": len(numbers),
    }


# ---------------------------------------------------------------------------
# SMILES → AtomsData
# ---------------------------------------------------------------------------


def smiles_to_atomsdata_core(smiles: str, seed: int = 2025) -> AtomsData:
    """Convert a SMILES string to an :class:`~chemgraph.schemas.atomsdata.AtomsData`.

    Parameters
    ----------
    smiles : str
        SMILES string representation of the molecule.
    seed : int, optional
        Random seed for RDKit 3D structure generation, by default 2025.

    Returns
    -------
    AtomsData
        Structure with no periodic boundary conditions.

    Raises
    ------
    ValueError
        If the SMILES string is invalid or 3D generation fails.
    """
    numbers, positions = smiles_to_3d(smiles, seed=seed)
    return AtomsData(
        numbers=numbers,
        positions=positions,
        cell=[[0, 0, 0], [0, 0, 0], [0, 0, 0]],
        pbc=[False, False, False],
    )