import logging as lg
from concurrent.futures import ProcessPoolExecutor
from functools import reduce
from typing import Callable, Iterable, cast

from rdkit import rdBase
from rdkit.Chem import (
    Kekulize,
    Mol,
    MolFromInchi,
    MolFromSmiles,
    MolToInchi,
    MolToInchiAndAuxInfo,
    MolToInchiKey,
    MolToSmiles,
    RemoveStereochemistry,
)
from rdkit.Chem.inchi import InchiToInchiKey

from fairfetched.utils._track import track

from ._types import StrToMolFn

# def _fuse_pipeline(fns: tuple[Callable, ...]) -> Callable:
#     return reduce(lambda f, g: lambda m: (r := f(m)) and g(r), fns)


def fuse_pipeline(
    *fns: Callable[[Mol | str | None], Mol | tuple[str, str, str] | None],
) -> Callable:
    """Combines functions on mols into one function. Technically made for mol-to-mol|None functions,
    but could be started and finished with Any->Mol and Mol->Any functions, respectively

    Example usage:
        ```
        standardise = _fuse_pipeline(remove_stereochemistry,validate_inchi,standardise_papyrus)
        ```
    """
    return reduce(lambda f, g: lambda m: (r := f(m)) and g(r), tuple(fns))


mol_to_smiles = cast(StrToMolFn, MolToSmiles)
mol_to_inchi = cast(StrToMolFn, MolToInchi)


def remove_stereochemistry_and_return_mol(mol: Mol) -> Mol:
    """wrapper returning mol after calling RemoveStereochemistry which returns None.

    Original docstring:
    ---
        Removes all stereochemistry info from the molecule.
        C++ signature :
            void RemoveStereochemistry(RDKit::ROMol {lvalue})
    """
    RemoveStereochemistry(mol)
    return mol


def validate_inchi(mol: Mol) -> Mol:
    """converts to and from inchi to ensure inchis are valid by running
    `rdkit.Chem.MolFromInchi(rdkit.Chem.MolToInchi(mol))`
    """
    return MolFromInchi(MolToInchi(mol))


def _mol_to_kekulised_smiles(mol: Mol) -> str:
    return MolToSmiles(mol, isomericSmiles=False, kekuleSmiles=True)


def _safe_mol_to_kekulised_smiles(mol: Mol) -> str:
    """Returns a SMILES string suitable for re-loading the molecule later.
    Removes stereochemistry.
    """
    try:
        Kekulize(mol, clearAromaticFlags=True)
    except Exception as e:
        lg.debug(f"Kekulization failed with error {e}")
        return MolToSmiles(mol, isomericSmiles=False)
    return MolToSmiles(mol, isomericSmiles=False)


def mol_from_smiles(mol):
    return MolFromSmiles(mol)


# def _safe_remove_stereochemistry(mol: Mol | None) -> Mol | None:
#     try:
#         RemoveStereochemistry(mol)
#     except Exception as e:
#         lg.error(f"Failed in RemoveStereochemistry(mol) with error {e}")
#         return None
#     return mol


# def _safe_validate_inchi(mol) -> Mol | None:
#     try:
#         inchi = MolToInchi(mol)
#         mol = MolFromInchi(inchi)
#         if not isinstance(mol, Mol):
#             lg.error(f"Failed on InChI conversion: {inchi}")
#             return None
#         return mol
#     finally:
#         return None


# def _safe_standardise_papyrus(mol: Mol | None, **kwargs) -> Mol | None:
#     if not mol:
#         return None
#     try:
#         mol, message = papyrus_standardise(
#             mol, return_type=True, raise_error=False, **kwargs
#         )
#         if not mol:
#             lg.error(f"Papyrus-type standardisation did not return a mol: {message}")
#             return None
#         return mol
#     except Exception as e:
#         lg.error(f"failed papyrus standardisation with exception: {e}")
#         return None


# def _safe_smiles_to_mol(smiles: str) -> Mol | None:
#     try:
#         return MolFromSmiles(smiles)
#     except Exception as e:
#         lg.error(f"Failed on `MolFromSmiles` conversion: {smiles} with exception {e}")
#         return None


# def _safe_inchi_to_mol(inchi: str) -> Mol | None:
#     try:
#         return MolFromInchi(inchi)
#     except Exception as e:
#         lg.error(e)
#         return None


# def mol_to_strings(mol: Mol | None) -> tuple[str, str, str, str] | None:
#     """returns a tuple of kekulised smiles, inchi, inchiauxinfo, and inchikeys"""
#     if not Mol:
#         return None
#     return MolToSmiles(mol), *MolToInchiAndAuxInfo(mol), MolToInchiKey(mol)


# standardise_aeffect = cast(
#     Callable[[Mol | None], Mol | None],
#     fuse_pipeline(
#         _safe_remove_stereochemistry,
#         partial(
#             _safe_standardise_papyrus,
#             filter_mixtures=True,  # @TODO: check if this is desired
#             filter_inorganic=False,
#             filter_non_small_molecule=False,
#             return_type=True,
#             raise_error=False,
#         ),
#         _safe_validate_inchi,
#     ),
# )


# # def _smiles_to_clean_mol(smiles: str) -> Mol | None:
# #     """Returns cleaned molecules from SMILES string or None upon failure or invalid molecules"""
# #     assert isinstance(smiles, str)
# #     if not smiles:
# #         lg.warning("empty smiles")
# #         return None
# #     mol = MolFromSmiles(smiles)
# #     if not isinstance(mol, Mol):
# #         lg.debug("Failed on initial `MolFromSmiles` conversion")
# #         return None
# #     try:
# #         RemoveStereochemistry(mol)
# #     except Exception as e:
# #         lg.error(f"Failed in RemoveStereochemistry(mol) with error {e}")
# #         return None
# #     # @TODO: check if stereochemistry removal should be before or after standardisation
# #     if not mol:
# #         lg.error(f"{smiles} could not be converted to a mol")
# #         return None
# #     try:
# #         mol, message = papyrus_standardise(
# #             mol,
# #             filter_mixtures=True,  # @TODO: check if this is desired
# #             filter_inorganic=False,
# #             filter_non_small_molecule=False,
# #             return_type=True,
# #             raise_error=False,
# #         )
# #         if not mol:
# #             lg.warning(
# #                 f"{smiles} could not be cleaned during Papyrus-type standardisation: {message}"
# #             )
# #             return None
# #         mol = MolFromInchi(MolToInchi(mol))
# #         if not mol:
# #             lg.warning(f"{smiles} did not give valid inchi: {message}")
# #         return mol
# #     except Exception as e:
# #         lg.error(
# #             f"{smiles} could not be cleaned in papyrus-style standardsation with exception: {e}"
# #         )
# #         return None
# #     return None


# def inchis_to_mols(
#     inchis: Iterable[str], cache: Path | None = None
# ) -> list[Mol | None]:
#     """Convert InChIs to RDKit Mol objects"""
#     return [
#         _safe_inchi_to_mol(i) for i in track(inchis, desc="Converting InChI to mol")
#     ]


# def mols_to_kekulised_smiles(mols: Iterable[Mol]) -> list[str]:
#     """Returns kekulised SMILES strings suitable for re-loading the molecule later.
#     Removes stereochemistry.
#     """
#     return [
#         _mol_to_kekulised_smiles(m)
#         for m in track(mols, desc="converting mols to kekulised smiles")
#     ]


# def _mol_to_inchi_and_auxinfo(mol: Mol) -> tuple[str, str]:
#     return MolToInchiAndAuxInfo(mol)


# def _mol_to_inchi(mol: Mol) -> str:
#     return MolToInchiAndAuxInfo(mol)[0]  # to maintain consistency


# def mols_to_inchis(mols: Iterable[Mol]) -> list[str]:
#     return [_mol_to_inchi(m) for m in track(mols, desc="converting mols to inchis")]


# def mols_to_inchis_and_auxinfo(mols: Iterable[Mol]) -> list[tuple[str, str]]:
#     return [
#         _mol_to_inchi_and_auxinfo(m)
#         for m in track(mols, desc="converting mols to inchis with auxinfo")
#     ]


# def inchis_to_inchikeys(inchis: Iterable[str]) -> list[str]:
#     return [
#         InchiToInchiKey(i) for i in track(inchis, desc="converting inchis to inchikeys")
#     ]


# def mols_to_inchikeys(mols: Iterable[Mol]) -> list[str]:
#     return [
#         InchiToInchiKey(MolToInchiAndAuxInfo(m)[0])
#         for m in track(mols, desc="converting mols to inchikeys")
#     ]


def mols_to_inchikeys(mols: Iterable[Mol]) -> list[str]:
    return [MolToInchiKey(m) for m in track(mols, desc="converting mols to inchikeys")]


# def mols_to_strings_zip(mols: Iterable[Mol]) -> zip[tuple[str, str, str, str]]:
#     return zip(
#         map(_mol_to_kekulised_smiles, mols),
#         *(zip(*map(_mol_to_inchi_and_auxinfo, mols))),
#         map(MolToInchiKey, mols),
#     )


def mols_to_strings_pooled(mols: Iterable[Mol]) -> tuple[tuple[str, str, str, str]]:
    with ProcessPoolExecutor() as p:
        smiles = list(p.map(_mol_to_kekulised_smiles, mols))
        i_a = list(p.map(MolToInchiAndAuxInfo, mols))
        inchikeys = list(p.map(MolToInchiKey, (inchi for inchi, aux in i_a)))

    return tuple(zip(smiles, *(zip(*i_a)), inchikeys))


def mols_to_strings(mols: Iterable[Mol]) -> list[tuple[str, str, str, str]]:
    return list(mols_to_strings(mols))


mol_to_kekulized_smiles = _mol_to_kekulised_smiles  # American spelling alias
inchi_to_inchikey = InchiToInchiKey  # convenience alias


# def make_rdkit_fingerprint_fn(
#     fp_type: str = "morgan", radius: int = 2, n_bits: int = 2048
# ):
#     """Returns a rep_fn that produces a bit vector as a list[int]."""

#     def fp_fn(mol) -> list[int] | None:
#         from rdkit.Chem import rdMolDescriptors

#         if fp_type == "morgan":
#             bv = rdMolDescriptors.GetMorganFingerprintAsBitVect(
#                 mol, radius, nBits=n_bits
#             )
#         elif fp_type == "rdkit":
#             from rdkit.Chem import RDKFingerprint

#             bv = RDKFingerprint(mol, fpSize=n_bits)
#         else:
#             raise ValueError(f"Unknown fp_type: {fp_type}")
#         return list(bv.ToList())

#     return fp_fn
