from functools import partial
from typing import TYPE_CHECKING, Callable, cast

from ._optional import _papyrus_standardize, chembl_get_parent_mol, chembl_standardize

if TYPE_CHECKING:
    from rdkit.Chem import Mol


def papyrus_standardize(mol: Mol) -> Mol | None:
    """Standardize a molecule from either a RDKit molecule or SMILES.

    Steps:
    1) Use the ChEMBL structure pipeline (get parent molecule & standardization)
    2) remove supplementary salts & metals
    3) check if a mixture
    4) check if organic
    5) uncharge
    5) check if 200 <= MW <= 800
    6) canonicalize tautomer
    7) Use the ChEMBL structure pipeline (get parent molecule & standardization)

    :param mol: RDKit molecule to be standardized
    :return: a tuple of either the standardized molecule or None and a standardization_result flag
    """
    return _papyrus_standardize(mol, return_type=False, raise_error=True)  # ty: ignore[invalid-return-type] #ty:ignore [unused-ignore-comment] #ty:ignore [unused-ignore-comment]


def get_papyrus_standardize_with_custom_kwargs(
    remove_additional_salts: bool = True,
    remove_additional_metals: bool = True,
    filter_mixtures: bool = True,
    filter_inorganic: bool = True,
    uncharge: bool = True,
    filter_non_small_molecule: bool = True,
    canonicalize_tautomer: bool = True,
    small_molecule_min_mw: int | float = 200,
    small_molecule_max_mw: int | float = 800,
    tautomer_allow_stereo_removal: bool = True,
    tautomer_max_tautomers: int = 4294967295,
) -> Callable[[Mol], Mol | None]:
    """build a papyrus standardisation function to pass to `fairfetched.standardize()`'s
    `std_fns` kwarg.
    This is merely a convenience function: for default settings, you can simply
    pass papyrus_structure_pipeline.standardize, or you can specify the function
    by yourself using
    ```
    functools.partial(
        papyrus_structure_pipeline.standardize,
        **kwargs
        )
    ```


    Returns:
        a Callable[[Mol],Mol] taking mols and returning mols or None if not valid.

    Raises:
        variously on errors, which is handled by fairfetched.standardize()
    Note:
        - `return_type=False` and `raise_error=True` are fixed to work with pipelines
            if you want them to work, you should use the original function
            `papyrus_structure_pipeline.standardize` directly, for example with
            ```
            df.with_columns(pl.col('mol').map_elements(
                papyrus_structure_pipeline.standardize,
                return_dtype=<proper_return_type_in_polars_format>
                **standardize_kwargs,
            ))
            ```
    """

    return cast(
        Callable[[Mol], Mol | None],
        partial(
            papyrus_standardize,
            remove_additional_salts=remove_additional_salts,
            remove_additional_metals=remove_additional_metals,
            filter_mixtures=filter_mixtures,
            filter_inorganic=filter_inorganic,
            uncharge=uncharge,
            filter_non_small_molecule=filter_non_small_molecule,
            canonicalize_tautomer=canonicalize_tautomer,
            small_molecule_min_mw=small_molecule_min_mw,
            small_molecule_max_mw=small_molecule_max_mw,
            tautomer_allow_stereo_removal=tautomer_allow_stereo_removal,
            tautomer_max_tautomers=tautomer_max_tautomers,
            return_type=False,  # these should be set for the pipeline to work
            raise_error=True,
        ),
    )


def chembl_standardize_to_parent_mol(mol: Mol) -> Mol | None:
    """standardises molecule according to ChEMBL standardisation,
    then gets the parent mol. does the following steps
    (as taken from `chembl_structure_pipeline` v1.2.4).

    (Uses installed `chembl_structure_pipeline`, so actual steps may vary
    in case `chembl_structure_pipeline` implementations change(d))


    function                      \tinferred equivalent in paper
    ---                           \t ---
    1) update_mol_valences(m)     \t"Standardise unknown stereochemistry"
    2) remove_sgroups_from_mol(m) \t"Clear S Group data from the molfile"
    3) kekulize_mol(m)            \t"Generate a kekulé form of the structure"
    4) remove_hs_from_mol(m)      \t"Remove explicit H atoms except some cases"
    5) normalize_mol(m)           \t"Normalise structure"
    6) uncharge_mol(m)            \t"Ensure molecule is neutralised, if possible"
    7) flatten_tartrate_mol(m)    \t"Normalise (straighten) triple bonds and allenes"
    8) cleanup_drawing_mol(m)     \t"Normalise (straighten) triple bonds and allenes"
    9) Chem.SanitizeMol(m)        \t"RDKit sanitation"
    10) get_parent_molecule       \t "All information about isotopes is removed,
                                  \t as are solvents and salts present in the
                                  \tmolfile which match any of the components
                                  \tin the defined salt and solvent lists.
                                  \tHaving removed all salts (e.g. Na+ that
                                  \tmight be included to neutralise a carboxyl
                                  \tgroup), the resulting molecule is
                                  \tneutralised and a new molfile created as
                                  \tthe “parent” molecule" (see article below)

    Note:
    ---
        (skips molecules that are 'excluded' molecules, see
        https://link.springer.com/article/10.1186/s13321-020-00456-1#Sec2
        (i.e. more than 7 borons, or any of [Sc], [Ti], [V],
        [Cr], [Mn], [Fe], [Co], [Ni], [Cu], [Ga], [Y], [Zr], [Nb], [Mo], [Tc],
        [Ru], [Rh], [Pd], [Cd], [In], [Sn], [La], [Hf], [Ta], [W], [Re], [Os],
        [Ir], [Pt], [Au], [Hg], [Tl], [Pb], [Bi], [Po], [Ac], [Ce], [Pr], [Nd],
        [Pm], [Sm], [Eu], [Gd], [Tb], [Dy], [Ho], [Er], [Tm], [Yb], [Lu], [Th],
        [Pa], [U], [Np], [Pu], [Am], [Cm], [Bk], [Cf], [Es], [Fm], [Md], [No],
        [Lr], [Ge], [Sb]))

    DOI:
    ---
       https://doi.org/10.1186/s13321-020-00456-1

    """
    return chembl_get_parent_mol(chembl_standardize(mol))


__all__ = [
    "papyrus_standardize",
    "get_papyrus_standardize_with_custom_kwargs",
    "chembl_standardize",
    "chembl_get_parent_mol",
    "chembl_standardize_to_parent_mol",
]
