from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING, Optional

import chemrecon
from chemrecon import ChemReaction
from chemrecon.chem.mol import MolInstance, MolTemplate
from mntools import Side

from aam_annotator.matching_algorithm import CannotMatch, matching_alg
from aam_annotator.problem import (
    CuratedStructureNotInAAM,
    MatchingProblem,
    MultiplicityMismatch,
    UnmatchedCuratedStructure,
)
from aam_annotator.utils.hydrogen import is_hydrogen

if TYPE_CHECKING:
    from aam_annotator.mapped_network import MappedMetabolite, MappedReaction


class Matching:
    """A matching between the metabolites of a reaction in the network model and the structures in an AAM.
    Note that matchings do not match hydrogens, and hydrons are explicitly ignored.
    """

    # TODO metabolite should be matched to same structure on both sides?
    #  - Add a check to ensure this?
    #  - or: redesign the data structure.

    reaction: MappedReaction
    aam: ChemReaction

    # Reaction and AAM info
    structures: dict[Side, list[MolTemplate]]
    structure_instances: dict[Side, list[MolInstance]]
    metabolites: dict[Side, list[MappedMetabolite]]

    # Matching
    match: dict[Side, dict[MappedMetabolite, MolTemplate]]
    match_conf: dict[Side, dict[tuple[MappedMetabolite, MolTemplate], float]]

    unmatched_metabolites: dict[Side, list[MappedMetabolite]]
    unmatched_structures: dict[Side, list[MolTemplate]]

    # Multiplicity of metabolites and structure templates
    metabolite_multiplicity: dict[Side, dict[MappedMetabolite, int]]
    structure_multiplicity: dict[Side, dict[MolTemplate, int]]

    # Misc
    template_instance_map: dict[MolTemplate, list[MolInstance]]
    instance_template_map: dict[MolInstance, MolTemplate]
    problems: list[MatchingProblem]

    # Getters
    # ---------------------------------------------------------------------------------------------------------
    def is_complete(self) -> bool:
        return (
            len(self.unmatched_metabolites[Side.L]) == 0
            and len(self.unmatched_metabolites[Side.R]) == 0
            and len(self.unmatched_structures[Side.L]) == 0
            and len(self.unmatched_structures[Side.R]) == 0
        )

    def get_conf(self) -> float:
        """Return avg. conf (counting unmatched items as 0)."""
        divisor = (
            min(
                max(
                    len(self.structures[Side.L]),
                    len(self.metabolites[Side.L]),
                ),
                max(
                    len(self.structures[Side.R]),
                    len(self.metabolites[Side.R]),
                ),
                1,
            )
            * 2
        )
        total_conf: float = 0.0
        for side in Side:
            for m, moltemp in self.match[side].items():
                total_conf += self.match_conf[side][(m, moltemp)]
        return total_conf / divisor

    def get_structure(self, metabolite: MappedMetabolite) -> Optional[MolTemplate]:
        """Return the structure that matches the given metabolite."""
        struct_l = self.match[Side.L].get(metabolite, None)
        struct_r = self.match[Side.R].get(metabolite, None)
        match struct_l, struct_r:
            case None, None:
                return None
            case s, None:
                return s
            case None, s:
                return s
            case s_l, s_r:
                if s_l != s_r:
                    raise ValueError(
                        "Mismatched structures for metabolite (different on side L than R: "
                        + str(metabolite)
                    )
                else:
                    return s_l

    def has_multiplicity_mismatch_on_side(self, side: Side) -> bool:
        """Return True if any multiplicities do not match."""
        metabolite_mult_counts: dict[int, int] = defaultdict(int)
        structure_mult_counts: dict[int, int] = defaultdict(int)
        for _, mult in self.metabolite_multiplicity[side].items():
            metabolite_mult_counts[mult] += 1
        for _, mult in self.structure_multiplicity[side].items():
            structure_mult_counts[mult] += 1

        return metabolite_mult_counts != structure_mult_counts

    def has_multiplicity_mismatch(self) -> bool:
        """Return True if any multiplicities do not match."""
        return self.has_multiplicity_mismatch_on_side(
            Side.L
        ) or self.has_multiplicity_mismatch_on_side(Side.R)

    # Creation
    # ----------------------------------------------------------------------------------------------------------
    def __init__(
        self,
        reaction: MappedReaction,
        aam: ChemReaction,
    ):
        """Initialize empty matching."""
        self.reaction = reaction
        self.aam = aam
        self.problems = list()

        # Initialize metabolites and structures
        self.structures = {Side.L: [], Side.R: []}
        self.structure_instances = {Side.L: [], Side.R: []}
        self.metabolites = {Side.L: [], Side.R: []}
        self.template_instance_map = {}
        self.instance_template_map = {}
        self.metabolite_multiplicity = {Side.L: {}, Side.R: {}}
        self.structure_multiplicity = {Side.L: {}, Side.R: {}}
        self.match = {Side.L: {}, Side.R: {}}
        self.match_conf = {Side.L: {}, Side.R: {}}
        self.unmatched_metabolites = {Side.L: [], Side.R: []}
        self.unmatched_structures = {Side.L: [], Side.R: []}

        # First, check if this is an all-hydrogen reaction.
        all_hydrogen_metabolites: dict[Side, bool] = {
            Side.L: all(m.is_hydrogen for m in reaction.get_lhs_metabolites().keys()),
            Side.R: all(m.is_hydrogen for m in reaction.get_rhs_metabolites().keys()),
        }
        all_hydrogen_structures: dict[Side, bool] = {
            Side.L: all(is_hydrogen(molinst) for molinst in aam.lhs_instances),
            Side.R: all(is_hydrogen(molinst) for molinst in aam.rhs_instances),
        }
        is_all_hydrogen_reaction = (
            all_hydrogen_metabolites[Side.L]
            and all_hydrogen_metabolites[Side.R]
            and all_hydrogen_structures[Side.L]
            and all_hydrogen_structures[Side.R]
        )
        if is_all_hydrogen_reaction:
            print("    - All hydrogen matching. Skipping.")
            # TODO skip
            pass
        else:
            # If all hydrogen on one side only, raise an error.
            if all_hydrogen_metabolites[Side.L] or all_hydrogen_metabolites[Side.R]:
                # TOOD hydrogen metabolites but not structures
                pass
            if all_hydrogen_structures[Side.L] or all_hydrogen_structures[Side.R]:
                # TODO hydrogen structures but not metabolites
                pass

        # Load metabolites and stuctures
        for side, model_metabolites_with_stoich, aam_structures in [
            (Side.L, reaction.get_lhs_metabolites(), aam.lhs_instances),
            (Side.R, reaction.get_rhs_metabolites(), aam.rhs_instances),
        ]:
            for m, stoich in model_metabolites_with_stoich.items():
                if TYPE_CHECKING:
                    assert isinstance(m, MappedMetabolite)
                if m.is_hydrogen and not is_all_hydrogen_reaction:
                    # Skip if hydrogen
                    continue
                self.metabolites[side].append(m)
                self.metabolite_multiplicity[side][m] = abs(stoich)

            # Gather templates
            temp_inst_map_side: dict[MolTemplate, list[MolInstance]] = defaultdict(list)
            for molinst in aam_structures:
                if is_hydrogen(molinst) and not is_all_hydrogen_reaction:
                    # Skip if hydrogen
                    continue
                mol_temp = molinst.to_mol_template()
                temp_inst_map_side[mol_temp].append(molinst)

            # Add templates
            for moltemp, molinst_list in temp_inst_map_side.items():
                self.structures[side].append(moltemp)
                self.structure_multiplicity[side][moltemp] = len(molinst_list)
                self.structure_instances[side].extend(molinst_list)

                if moltemp not in self.template_instance_map:
                    self.template_instance_map[moltemp] = molinst_list
                else:
                    self.template_instance_map[moltemp].extend(molinst_list)

                for molinst in molinst_list:
                    self.instance_template_map[molinst] = moltemp

            # Set all metabolites and structures to be unmatched
            self.unmatched_metabolites[side] = self.metabolites[side].copy()
            self.unmatched_structures[side] = self.structures[side].copy()

        # Check whether multiplicities match
        # TODO ?

        # Done with init.
        pass

    # Add matchings
    # ----------------------------------------------------------------------------------------------------------
    def match_metabolite(
        self,
        side: Side,
        metabolite: MappedMetabolite,
        mol: MolTemplate,
        conf: float = 1.0,
    ):
        """Produdce a matching."""
        # Check (inclusion)
        if metabolite not in self.metabolites[side]:
            raise ValueError(f"Metabolite {metabolite} not found on side {side}.")
        if mol not in self.structures[side]:
            raise ValueError(f"Structure {mol} not found on side {side}.")

        if metabolite not in self.unmatched_metabolites[side]:
            raise ValueError(f"Metabolite {metabolite} already matched.")
        if mol not in self.unmatched_structures[side]:
            raise ValueError(f"Structure {mol} already matched.")

        # Checks (multiplicity)
        mult_metabolite = self.metabolite_multiplicity[side][metabolite]
        mult_structure = self.structure_multiplicity[side][mol]
        if mult_metabolite != mult_structure:
            self.add_problem(MultiplicityMismatch(metabolite, mol))

        # Match
        self.match[side][metabolite] = mol
        self.match_conf[side][(metabolite, mol)] = conf

        self.unmatched_metabolites[side].remove(metabolite)
        self.unmatched_structures[side].remove(mol)

    # Other
    # ----------------------------------------------------------------------------------------------------------
    # Post-process after adding all matches
    def process(self):
        """After adding all matches, process to check self-consistency etc."""
        pass

    # Problems
    def add_problem(self, problem: MatchingProblem):
        self.problems.append(problem)

    # Matching alg
    # ----------------------------------------------------------------------------------------------------------
    def match_from_curations(self):
        """Perform matchings based on given curations."""
        # TODO check for match up to FICTS.
        for side in Side:
            for m in self.unmatched_metabolites[side]:
                if m.curated_mol is not None:
                    # Check for match on the same side
                    if m.curated_mol in self.unmatched_structures[side]:
                        # Add this matching.
                        try:
                            self.match_metabolite(side, m, m.curated_mol)
                        except ValueError:
                            self.add_problem(CuratedStructureNotInAAM(m, m.curated_mol))
                    else:
                        # No corresponding mol
                        self.add_problem(UnmatchedCuratedStructure(m, m.curated_mol))

    def match_structures(self):
        """Main matching algorithm. Will auto-match unmatched metabolites and structures based on information
        in the network.
        """
        matches: dict[Side, list[tuple[MappedMetabolite, MolTemplate]]] = defaultdict(
            list
        )
        conf: dict[Side, float] = dict()

        for side in Side:
            match_metabolites: list[MappedMetabolite] = list()
            for m in self.unmatched_metabolites[side]:
                for _ in range(self.metabolite_multiplicity[side][m]):
                    match_metabolites.append(m)

            match_structures: list[MolTemplate] = list()
            for struct in self.unmatched_structures[side]:
                for _ in range(self.structure_multiplicity[side][struct]):
                    match_structures.append(struct)

            try:
                match_side, conf_side = matching_alg(
                    match_metabolites,
                    match_structures,
                    consider_n_most_confident_strucutres=3,
                )
                conf[side] = conf_side

                # Add matches
                for m, moltemp in match_side:
                    matches[side].append((m, moltemp))

            except CannotMatch as exc:
                # Could not perform matching on this side.
                self.add_problem(exc)
                continue

        # Add matchings
        for side, match_side in matches.items():
            for m, moltemp in match_side:
                try:
                    self.match_metabolite(side, m, moltemp, conf=conf[side])
                except ValueError:
                    # Could not assign match
                    # Likely because the metabolite was already matched to a different structure
                    if self.has_multiplicity_mismatch_on_side(side):
                        # Expected
                        continue
                    else:
                        # Unexpected
                        continue


# Different types of matching
# ==============================================================================================================
class DatabaseAAMMatching(Matching):
    """A matching based on a manually curated AAM."""

    aam_entry: chemrecon.AAM
    aam_conf: float
    match_conf: float

    def __init__(
        self,
        reaction: MappedReaction,
        aam: ChemReaction,
        aam_entry: chemrecon.AAM,
        aam_conf: float,
    ):
        super().__init__(reaction, aam)
        self.aam_entry = aam_entry
        self.aam_conf = aam_conf


class TrivialTransportMatching(Matching):
    """A matching based on a trivial transport reaction."""


class GeneratedMatching(Matching):
    """A matching based on a generated reaction."""


class UserCuratedAAMMatching(Matching):
    """A matching based on a user-curated AAM."""


# Generating matchings
# ==============================================================================================================
