""" Defines a wrapper class for RDKit molecules.
    A Mol may either be a MolTemplate or a MolInstance.
    - MolInstance: The molecule as it appears in a reaction concretely.
    - MolTemplate: More abstract representation of the molecule in general.
"""
from __future__ import annotations

import ast
from copy import copy, deepcopy
from typing import Type, Optional, Any

import rdkit.Chem as rdk
from rdkit.Chem import rdFingerprintGenerator as rdk_fp
import rdkit.Chem.MolStandardize.rdMolStandardize as rdk_std
from rdkit import DataStructs as rdk_ds

from chemrecon.chem.elements import atomicnum_element
from chemrecon.chem.sumformula import SumFormula
from chemrecon.schema import MolStructure
from chemrecon.schema.enums import FeatureEnum

# RDkit generators
fpgen = rdk_fp.GetRDKitFPGenerator()

# Main class
# --------------------------------------------------------------------------------------------------------------
class Mol:
    """ Wrapper for RDKit Mol
    """
    # TODO rename/refactor to 'MolStructure'?
    mol:                rdk.Mol

    smiles:             Optional[str]
    features:           set[Type[Feature]]  # Features for which the rdk_mol is standardized
    n_atoms:            int
    mass:               float
    charge:             int
    molformula:         SumFormula
    provenance:         Optional[str]
    propdict:           dict[str, str]  # Properties



    # Optional fingerprint
    fp:                 Optional[Any]

    def __init__(
        self,
        rdk_mol: rdk.Mol,
        set_features: set[Type[Feature]] = None,
        provenance: Optional[str] = None
    ):
        if rdk_mol is None:
            raise ValueError('Cannot create Mol from None.')

        # Normalize the molecule
        try:
            self.mol = rdk_std.Normalize(rdk_mol)       # TODO what exactly does normalisation do?
        except ValueError as e:
            # Could not create molecule, or molecule is None?
            print(e)
            self.mol = None
            self.smiles = None
            # return
            raise ValueError(f'Could not create molecule: {e}')

        # Set smiles for printing
        self.smiles = rdk.MolToSmiles(self.mol)

        # Fix SMILES bug (non-C lowercase atom symbols)
        # https://github.com/rdkit/rdkit/issues/3697
        self.smiles = self.smiles.replace('[o', '[O')

        if self.smiles is None:
            raise ValueError('TODO fix')

        # Compute featuers
        if set_features is not None:
            # Features set in constructor
            self.features = set_features
        else:
            self.features = set()
            for f in feats:
                # Check standardization and set
                if f.is_standardized(mol = self):
                    self.features.add(f)

        # Get properties from name if initialized through an RXN file
        try:
            molprops_rxn = rdk_mol.GetProp('_Name')
            try:
                rxn_propdict = ast.literal_eval(molprops_rxn)
                self.propdict = rxn_propdict
            except Exception as e:
                self.propdict = dict()
        except KeyError:
            # Not applicable, skip
            self.propdict = dict()
            pass

        # Populate misc fields
        self.n_atoms = self.mol.GetNumAtoms(onlyExplicit = False)
        self.provenance = provenance

        # Get MolFormula
        self.molformula = self.get_molformula()

        # TODO n_atoms_tracked
        # TODO mass (simple)

        self.charge = rdk.GetFormalCharge(self.mol)

        # Set fp
        self.fp = None

        if self.provenance is None:
            # TODO fix
            pass

    # Converters
    # ----------------------------------------------------------------------------------------------------------
    def to_smiles(self) -> str:
        # TODO return more details
        return self.smiles

    # Database interfacing
    # ----------------------------------------------------------------------------------------------------------
    def to_database_struct(self) -> MolStructure:
        """ Create a database row to be inserted."""
        if self.mol is None and self.smiles is None:
            raise ValueError('Cannot convert None entry')

        # TODO fixes rdkit bug?
        # Fix SMILES bug (non-C lowercase atom symbols)
        # https://github.com/rdkit/rdkit/issues/3697
        self.smiles.replace(' [o', '[O')

        # Else, convert
        return MolStructure(
            smiles = self.to_smiles(),
            std_feats = [
                f.feature_enum for f in self.features
            ]
        )

        # TODO insert calculated properties

    # Calculate properties
    # ----------------------------------------------------------------------------------------------------------
    # TODO

    # Standardisation
    def get_standardised(self, feat_list: Optional[list[Feature]] = None) -> Mol:
        """ Get a new Mol standardised according to the given features. If None given, standardise according
            to all features.
        """
        feat_list = feat_list or [F, I, C, T, S]
        rdk_s: rdk.Mol = self.mol
        for f in feat_list:
            rdk_s = f.standardise_rdk(rdk_s)
        return Mol(rdk_s, set_features = feat_list, provenance = self.provenance)

    # Identity and similarity
    # ------------------------------------------------------------------------------------------------------------------
    def is_identical_up_to_map(self, other: Mol):
        """ Checks whether two Mols are identical except for extra details such as atom mapping numbers. """
        m: rdk.Mol = deepcopy(self.mol)
        m_: rdk.Mol = deepcopy(other.mol)

        for a in m.GetAtoms():
            a.SetAtomMapNum(0)
        for a in m_.GetAtoms():
            a.SetAtomMapNum(0)

        return rdk.MolToSmiles(m) == rdk.MolToSmiles(m_)

    def is_identical_up_to(self, other: Mol, features: set[Feature]) -> bool:
        # TODO
        raise NotImplementedError

    def generate_fingerprint(self):
        if self.fp is None:
            self.fp = fpgen.GetFingerprint(self.mol)

    def get_similarity(self, other: Mol) -> float:
        """ Get a similarity metric between two Mols. Should be 1 if the mols are identical,
            and approach zero as the difference increases.
        """
        self.generate_fingerprint()
        other.generate_fingerprint()
        return rdk_ds.TanimotoSimilarity(
            self.fp, other.fp
        )

    # Misc
    # ----------------------------------------------------------------------------------------------------------
    def feature_string(self) -> str:
        s = ''
        for f in feats:
            if f in self.features:
                s += f.symbol
            else:
                s += '-'
        return s

    def __repr__(self):
        return f'{self.smiles}    [{self.feature_string()}]'

    def __hash__(self):
        return f'{self.smiles}:::{self.feature_string()}'.__hash__()

    def __eq__(self, other: Mol):
        if isinstance(other, Mol):
            return self.smiles == other.smiles and self.features == other.features
        else:
            return False

    # Serialise
    def serialize(self) -> dict:
        return {
            'smiles': self.to_smiles(),
            'features': self.feature_string(),
            'molformula': 'TODO'  # TODO
        }

    # Get molecular formula
    def get_molformula(self) -> SumFormula:
        """ Returns the molecular formula
        """
        # TODO requires rdkit documentation
        formula_num: dict[int, int] = dict()
        for a in self.mol.GetAtoms():
            n: int = a.GetAtomicNum()
            formula_num[n] = formula_num.get(n, 0) + 1

        charge = rdk.GetFormalCharge(self.mol)

        # Translate to formula of elements and return
        return SumFormula(
            formula = {
                atomicnum_element[n]: i for n, i in formula_num.items()
            },
            charge = charge
        )

# Mol Instance and Template
# ----------------------------------------------------------------------------------------------------------------------
class MolInstance(Mol):
    """ MolInstance represents a particular instance of a molecule as it participates in a reaction.
    """

    n_atoms_tracked: int
    _atom_map: Optional[list[int]]  # For each atom in the MolInstance, gives the mapping.
    _atom_map_smiles_order: Optional[list[int]]

    def __init__(self, rdk_mol: rdk.Mol, set_features: set[Type[Feature]] = None, provenance: Optional[str] = None):
        super().__init__(rdk_mol, set_features = set_features, provenance = provenance)
        self._atom_map = None
        self._atom_map_smiles_order = None

    def to_mol_template(self) -> MolTemplate:
        """ Remove all information of the instance as it participates in a reaction (atom-to-atom map etc.)
            to generate a generic template of the structure.
        """
        mtemp = MolTemplate(
            rdk_mol = self.mol,
            set_features = self.features,
            provenance = self.provenance
        )

        # Carry over certain properties
        mtemp.propdict = self.propdict
        return mtemp

    def get_atom_map_in_native_order(self) -> list[int]:
        # The numbering of the atoms as in the reaction it is an instance of.
        if not self._atom_map:
            self._atom_map = [
                a.GetAtomMapNum()
                for a in self.mol.GetAtoms()
            ]

        return self._atom_map


    def get_atom_map_in_smiles_order(self, safe: bool = False) -> list[int]:
        # Same as getting the atom map, but with the ordering of the atoms in the smiles string
        # Needed to populate DB with atom-to-atom maps
        if not self._atom_map_smiles_order:
            smilesmol = rdk.MolFromSmiles(self.smiles, sanitize = not safe)
            self._atom_map_smiles_order = [
                a.GetAtomMapNum()
                for a in smilesmol.GetAtoms()
            ]

        return self._atom_map_smiles_order

    def serialize(self) -> dict:
        d = super().serialize()
        d.update({
            'template': self.to_mol_template().to_smiles()
        })
        return d


class MolTemplate(Mol):
    """ MolTemplate represents a consistent type of molecule, independent of its particular participation and
        mapping in a reaction.
    """
    pass

    def __init__(
        self,
        rdk_mol: rdk.Mol,
        set_features: set[Type[Feature]] = None,
        provenance: Optional[str] = None
    ):
        rdk_mol_ = deepcopy(rdk_mol)
        if not rdk_mol_:
            raise ValueError(f'Could not create molecule')

        # Remove atom-to-atom maps
        for a in rdk_mol_.GetAtoms():
            a.SetAtomMapNum(0)

        super().__init__(rdk_mol = rdk_mol_, set_features = set_features, provenance = provenance)

    def instantiate(self, mapping: list[int]) -> MolInstance:
        """ Instantiate a template with a map, to produce an instance of the template as it participates in a reaction.
        """
        rdk_mol_ = deepcopy(self.mol)
        for a, mapnum in zip(rdk_mol_.GetAtoms(), mapping):
            a.SetAtomMapNum(mapnum)
        return MolInstance(rdk_mol = rdk_mol_, set_features = self.features, provenance = 'from_template')


# Features and standardization
# ----------------------------------------------------------------------------------------------------------------------
# ABC
class Feature:
    symbol: str
    name: str
    feature_enum: FeatureEnum

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        raise NotImplementedError

    @classmethod
    def standardise(cls, mol: Mol) -> rdk.Mol:
        if mol.mol is None:
            raise ValueError
        return cls.standardise_rdk(mol.mol)

    @classmethod
    def standardise_rdk(cls, rdk_mol: rdk.Mol) -> rdk.Mol:
        raise NotImplementedError


# FICTS features
# --------------------------------------------------------------------------------------------------------------

# [F] - Fragments
class F(Feature):
    symbol = 'F'
    name = 'Fragments'
    feature_enum = FeatureEnum.F

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        n_frags = rdk.GetMolFrags(mol.mol)
        return len(n_frags) <= 1

    @classmethod
    def standardise_rdk(cls, rdk_mol: rdk.Mol) -> rdk.Mol:
        mol_ = rdk_std.FragmentParent(rdk_mol, skipStandardize = True)
        return mol_


# [I] - Isotope
class I(Feature):
    symbol = 'I'
    name = 'Isotope'
    feature_enum = FeatureEnum.I

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        a: rdk.Atom
        for a in mol.mol.GetAtoms():
            if a.GetIsotope():
                return False
        return True

    @classmethod
    def standardise_rdk(self, rdk_mol: rdk.Mol) -> rdk.Mol:
        mol_ = copy(rdk_mol)
        a: rdk.Atom
        for a in mol_.GetAtoms():
            if a.GetIsotope():
                a.SetIsotope(0)
        return mol_


# [C] - Charge
class C(Feature):
    symbol = 'C'
    name = 'Charge'
    feature_enum = FeatureEnum.C

    # TODO use other uncharging approach?
    uncharger = rdk_std.Uncharger()

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        mol_ = cls.uncharger.uncharge(mol.mol)
        return mol_identical(mol.mol, mol_)

    @classmethod
    def standardise_rdk(cls, rdk_mol: rdk.Mol) -> rdk.Mol:
        mol_ = cls.uncharger.uncharge(rdk_mol)
        return mol_


# [T] - Tautomer
class T(Feature):
    symbol = 'T'
    name = 'Tautomer'
    feature_enum = FeatureEnum.T

    tautomer_enumerator = rdk_std.TautomerEnumerator()

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        """ Checks if the generated canonical tautomer is different. """
        try:
            mol_canon = cls.tautomer_enumerator.Canonicalize(mol.mol)
            return mol_identical(mol.mol, mol_canon)
        except rdk.rdchem.AtomKekulizeException:
            return False

    @classmethod
    def standardise_rdk(cls, rdk_mol: rdk.Mol) -> rdk.Mol:
        mol_ = cls.tautomer_enumerator.Canonicalize(rdk_mol)
        return mol_


# [S] - Stereo
class S(Feature):
    symbol = 'S'
    name = 'Stereo'
    feature_enum = FeatureEnum.S

    @classmethod
    def is_standardized(cls, mol: Mol) -> bool:
        # TODO redo this check
        return not ('/' in mol.smiles or '\\' in mol.smiles or '@' in mol.smiles)

    @classmethod
    def standardise_rdk(cls, rdk_mol: rdk.Mol) -> rdk.Mol:
        mol_ = copy(rdk_mol)
        rdk.RemoveStereochemistry(mol_)
        return mol_


feat_enum_map: dict[FeatureEnum, Type[Feature]] = {
    FeatureEnum.F: F,
    FeatureEnum.I: I,
    FeatureEnum.C: C,
    FeatureEnum.T: T,
    FeatureEnum.S: S,
}

# List of features
feats: list[Type[Feature]] = [F, I, C, T, S]


# Utils
# --------------------------------------------------------------------------------------------------------------
def mol_identical(a: rdk.Mol, b: rdk.Mol) -> bool:
    return a.HasSubstructMatch(b) and b.HasSubstructMatch(a)



