""" Implements a wrapper for the RDKit reaction type."""
from __future__ import annotations

from enum import Enum
from typing import Optional

import rdkit.Chem.rdChemReactions as rdk_r
from rdkit.Chem.rdChemReactions import ChemicalReaction
from rdkit.Chem import Draw as rdk_draw

from chemrecon.chem.mol import MolTemplate, MolInstance
from chemrecon.chem.sumformula import SumFormula
from chemrecon.schema.entry_types.aam import AAM


class Side(Enum):
    L = -1
    R = 1

type stoich = int

class ChemReaction:
    reaction: ChemicalReaction
    _reaction_smiles: Optional[str]

    # Templates represent the structure of each compound. (There can be duplicate templates)
    lhs_templates: dict[MolTemplate, stoich]
    rhs_templates: dict[MolTemplate, stoich]

    # Instances represent specific structures with maps as they participate in the reaction.
    lhs_instances: list[MolInstance]
    rhs_instances: list[MolInstance]

    # Mapping between each instance and its corresponding template
    instance_template_dict = dict[MolInstance, MolTemplate]

    # Properties of the given instances, as read through the input files if possible
    instance_properties: dict[MolInstance, dict[str, str]]
    template_ids: dict[MolTemplate, str]
    template_names: dict[MolTemplate, str]

    # Atom-to-atom map of this reaction
    map: dict[tuple[MolInstance, int], tuple[MolInstance, int]]
    lhs_index: dict[int, tuple[MolInstance, int]]  # The global map number to (molecule, index) tuple
    rhs_index: dict[int, tuple[MolInstance, int]]  # The global map number to (molecule, index) tuple

    def __init__(self, rdk_reaction: ChemicalReaction):
        self.reaction = rdk_reaction
        self._reaction_smiles = None

        self.lhs_instances = list()
        self.rhs_instances = list()
        self.lhs_templates = dict()
        self.rhs_templates = dict()

        self.instance_template_dict = dict()
        self.map = dict()

        self.instance_properties = dict()
        self.template_ids = dict()
        self.template_names = dict()

        # Set reactant and product MolInstances and map
        for rdk_mol in self.reaction.GetReactants():
            self.lhs_instances.append(
                MolInstance(rdk_mol, provenance = 'implicit')
            )

        for rdk_mol in self.reaction.GetProducts():
            self.rhs_instances.append(
                MolInstance(rdk_mol, provenance = 'implicit')
            )

        # Populate map
        # The global map number to (molecule, index) tuple
        # Global map number -> molecule, index (referring to the order of atoms in the SMILES string)
        self.lhs_index = dict()
        self.rhs_index = dict()
        for lhs_instance in self.lhs_instances:
            for local_index, global_index in enumerate(lhs_instance.get_atom_map_in_native_order()):
                self.lhs_index[global_index] = (lhs_instance, local_index)
        for rhs_instance in self.rhs_instances:
            for local_index, global_index in enumerate(rhs_instance.get_atom_map_in_native_order()):
                self.rhs_index[global_index] = (rhs_instance, local_index)

        self.map = dict()
        for i, (lhs_mol, lhs_index) in self.lhs_index.items():
            try:
                self.map[(lhs_mol, lhs_index)] = self.rhs_index[i]
            except KeyError:
                # Missing map
                # TODO do something if this is not a hydrogen?
                pass

        # self.map = {
        #     (lhs_mol, lhs_index): self.rhs_index[i]
        #     for i, (lhs_mol, lhs_index) in self.lhs_index.items()
        # }

        # Compute templates and add
        for side, instance_list, template_list in [
            (-1, self.lhs_instances, self.lhs_templates),
            (1, self.rhs_instances, self.rhs_templates)
        ]:
            mol_instance: MolInstance
            for mol_instance in instance_list:
                template = mol_instance.to_mol_template()

                # Add to instance-template dict
                self.instance_template_dict[mol_instance] = template

                # Add to template lists
                template_list[template] = template_list.get(template, 0) + side

    # Getters
    # ------------------------------------------------------------------------------------------------------------------
    def get_lhs_templates(self) -> list[MolTemplate]:
        return list(self.lhs_templates.keys())

    def get_rhs_templates(self) -> list[MolTemplate]:
        return list(self.rhs_templates.keys())

    # Balance
    # ------------------------------------------------------------------------------------------------------------------
    def get_balance_difference(self) -> SumFormula:
        """ Get the difference between the LHS and RHS as a (possible negative) MolFormula.
            Positive counts indicate a surplus on the LHS, negative counts indicate a surplus on the RHS.
        """
        raise NotImplementedError()

        # lhs_sum = sum(m.get_molformula() for m in self.lhs_instances)
        # rhs_sum = sum(m.get_molformula() for m in self.rhs_instances)
        # return lhs_sum - rhs_sum

    def is_balanced(self) -> bool:
        """ Returns true if balanced in both atomic composition and charge.
            If not balanced, use .get_balance_differenec() to inspect the difference between LHS and RHS.
        """
        return self.get_balance_difference().is_zero()

    # Representations
    # ------------------------------------------------------------------------------------------------------------------
    def to_reaction_smiles(self) -> str:
        if self._reaction_smiles is None:
            # self._reaction_smiles = rdk_r.ReactionToSmarts(self.reaction)
            self._reaction_smiles = rdk_r.ReactionToSmiles(self.reaction)

        return self._reaction_smiles

    # Serialise
    def serialize(self) -> dict:
        return {
            'reaction_smiles': self.to_reaction_smiles(),
            'lhs': [
                mol.serialize() for mol in self.lhs_instances
            ],
            'rhs': [
                mol.serialize() for mol in self.rhs_instances
            ],
            'balanced': 'TODO',  # TODO
            'balance_difference': 'TODO',  # TODO compute difference in sum formulae for LHS and RHS
        }

    # Misc
    # ------------------------------------------------------------------------------------------------------------------
    def __hash__(self):
        return self.to_reaction_smiles().__hash__()

    def sanity_check(self):
        """ Raises an exception if the AAM maps atoms of different elements.
        """
        # Check all instances have smiles
        for inst in [*self.lhs_instances, *self.rhs_instances]:
            if inst.smiles is None:
                raise AssertionError('Instance invalid.')

        # Check mapping
        for global_index, (molinst_l, local_index_l) in self.lhs_index.items():
            try:
                molinst_r, local_index_r = self.rhs_index[global_index]
            except KeyError:
                # Atom 'disappears'
                continue

            l_atom = molinst_l.mol.GetAtomWithIdx(local_index_l)
            r_atom = molinst_r.mol.GetAtomWithIdx(local_index_r)
            if l_atom.GetAtomicNum() != r_atom.GetAtomicNum():
                raise AssertionError(f'Element mismatch: L: {local_index_l}, R: {local_index_r}')




    # Visualisation
    # ------------------------------------------------------------------------------------------------------------------
    def show(self):
        img = rdk_draw.ReactionToImage(rxn = self.reaction, subImgSize = (800, 800))
        img.show()
        pass

# Creators
# ----------------------------------------------------------------------------------------------------------------------
def chem_reaction_from_aam_entry(
    entry: AAM
) -> ChemReaction:
    """ Given an AAM entry, load into a ChemReaction object.
    """
    return chem_reaction_from_reactionsmiles(
        entry.reaction_smiles
    )

def chem_reaction_from_reactionsmiles(
    reactionsmiles: str
) -> ChemReaction:
    rdk_reaction = rdk_r.ReactionFromSmarts(reactionsmiles)
    return ChemReaction(rdk_reaction)

def chem_reaction_from_rxn(
    rxn: str,
    provenance: Optional[str] = None,
    safe: bool = False
) -> ChemReaction:
    rdk_reaction = rdk_r.ReactionFromRxnBlock(rxn, sanitize = not safe, removeHs = not safe)
    return ChemReaction(rdk_reaction)
