"""Atom-atom maps and structural information on top of a metabolic network."""

from __future__ import annotations

from collections import defaultdict
from typing import Any, Iterable, Optional, OrderedDict, cast

import chemrecon
import mntools
from chemrecon.chem.chemreaction import ChemReaction
from chemrecon.chem.mol import Mol, MolTemplate
from mntools import Side

from aam_annotator.matching import Matching
from aam_annotator.problem import (
    DatabaseAAMProblem,
    MetaboliteProblem,
    Problem,
    ReactionProblem,
)
from aam_annotator.provenance import ProvenanceAAM, ProvenanceStruct
from aam_annotator.utils.std_id import (
    standardize_metabolite_id,
    standardize_reaction_id,
)


class MappedSpecies(mntools.Species):
    # Overrides
    # pyrefly: ignore [bad-override-mutable-attribute]
    associated_metabolite: Optional[MappedMetabolite]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _in_reactions: dict[MappedReaction, int]


class MappedMetabolite(mntools.Metabolite):
    # Curation
    curated_mol: Optional[MolTemplate]

    # Selection
    selected_mol: Optional[MolTemplate]
    selection_provenance: Optional[ProvenanceStruct]
    selection_details: str

    # Candidates
    candidate_structures_eg: Optional[
        OrderedDict[chemrecon.Mol, float]
    ]  # EG-based candidates
    candidate_structures_aam: Optional[
        dict[MolTemplate, list[tuple[MappedReaction, Matching]]]
    ]
    candidate_structures_aam_selected: Optional[
        dict[MolTemplate, list[tuple[MappedReaction, Matching]]]
    ]

    canonical_maps: dict[MolTemplate, ChemReaction]

    # Misc
    is_hydrogen: bool
    problems: list[MetaboliteProblem]

    # Overrides
    # pyrefly: ignore [bad-override-mutable-attribute]
    _species: list[MappedSpecies]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _species_by_compartment: dict[mntools.Compartment, MappedSpecies]

    def get_reactions_mapped(self) -> list[tuple[MappedReaction, mntools.StoichCoeff]]:
        """Get a list of reactions this metabolite participates in, with stoichiometric coefficient."""
        # TODO has been implemented upstream
        reactions: list[tuple[MappedReaction, mntools.StoichCoeff]] = []
        for s in self._species:
            for r, coeff in s._in_reactions.items():
                reactions.append((r, coeff))

        return reactions

    # Setters
    def set_curated_structure(self, mol: MolTemplate):
        if self.curated_mol is not None:
            raise ValueError("Curated structure already set!")
        self.curated_mol = mol
        self.set_selected_structure(mol, ProvenanceStruct.USER_CURATED)

    def set_selected_structure(self, mol: MolTemplate, provenance: ProvenanceStruct):
        if self.selected_mol is not None:
            raise ValueError("Selected structure already set!")
        self.selected_mol = mol
        self.selection_provenance = provenance

    def set_canonical_map_aam(self, alt_mol: MolTemplate, aam: ChemReaction):
        """Set a mapping from an alternative structure to the canonical structure."""
        # Checks
        if alt_mol in self.canonical_maps:
            raise ValueError(f"Canonical map for {alt_mol} already set!")
        assert self.candidate_structures_aam_selected is not None
        if alt_mol not in self.candidate_structures_aam_selected:
            raise ValueError(f"No candidate structure for {alt_mol}!")

        # Set
        self.canonical_maps[alt_mol] = aam

    def __init__(
        self, species: Iterable[mntools.Species], unified_id: str, unified_by: str
    ):
        super().__init__(species, unified_id, unified_by)

        # Initialize
        self.curated_mol = None
        self.selected_mol = None
        self.selection_provenance = None
        self.selection_details = ""

        self.candidate_structures_eg = None
        self.candidate_structures_aam = None
        self.candidate_structures_aam_selected = None

        self.canonical_maps = dict()

        self.problems = list()


class MappedReaction(mntools.Reaction):
    # Curation
    curated_aam: Optional[
        ChemReaction
    ]  # Curated AAM (AAM can be curated without matching?)
    curated_aam_matching: Optional[Matching]  # Curated matching, if given

    # Selection
    selected_aam_matching: Optional[Matching]
    selection_provenance: Optional[ProvenanceAAM]
    selection_details: str
    can_generate_reaction: bool

    # Candidates
    candidate_aams: list[tuple[Optional[chemrecon.AAM], ChemReaction, float]]
    candidate_matchings: list[Matching]

    # Problems
    problems: list[Problem]

    # Overrides
    # pyrefly: ignore [bad-override-mutable-attribute]
    _species: dict[mntools.Side, dict[MappedSpecies, float]]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _metabolites: dict[mntools.Side, set[MappedMetabolite]]
    _compartments: dict[mntools.Side, set[mntools.Compartment]]

    def __init__(
        self,
        id: str,
        name: Optional[str],
        lhs: dict[mntools.Species, mntools.StoichCoeff],
        rhs: dict[mntools.Species, mntools.StoichCoeff],
    ):
        super().__init__(id, name, lhs, rhs)
        self.candidate_aams = []
        self.problems = list()
        self.can_generate_reaction = False

    def get_neighbours_mapped(self) -> dict[MappedReaction, set[MappedMetabolite]]:
        """Get a list of 'neighbouring' reactions, that is, reactions which share at least one metabolite."""
        # TODO has been implemented upstream

        neighbours: dict[MappedReaction, set[MappedMetabolite]] = dict()
        for m in {*self._metabolites[Side.L], *self._metabolites[Side.R]}:
            m: MappedMetabolite
            for r, _ in m.get_reactions_mapped():
                if r not in neighbours:
                    neighbours[r] = set()
                neighbours[r].add(m)

        return neighbours

    # Setters
    def set_curated_aam(self, aam: ChemReaction):
        if self.curated_aam is not None:
            raise ValueError("Curated structure already set!")
        self.curated_aam = aam


class MappedNetwork(mntools.MetabolicNetwork):
    # Overrides
    # pyrefly: ignore [bad-override-mutable-attribute]
    _species: list[MappedSpecies]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _metabolites: list[MappedMetabolite]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _reactions: list[MappedReaction]

    # pyrefly: ignore [bad-override-mutable-attribute]
    _species_lookup: dict[str, MappedSpecies]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _metabolite_lookup: dict[str, MappedMetabolite]
    # pyrefly: ignore [bad-override-mutable-attribute]
    _reaction_lookup: dict[str, MappedReaction]

    # Standardized id lookup
    _species_std_id_lookup: dict[str, set[MappedSpecies]]
    _metabolite_std_id_lookup: dict[str, MappedMetabolite]
    _reaction_std_id_lookup: dict[str, MappedReaction]

    # Problems
    problems: list[Problem]

    # Setters
    # ----------------------------------------------------------------------------------------------------------
    def set_metabolite_curated_structure(
        self, metabolite: MappedMetabolite, mol: MolTemplate
    ):
        metabolite.set_curated_structure(mol)

    def set_candidate_structures_for_metabolite(
        self,
        metabolite: MappedMetabolite,
        candidate_structures: OrderedDict[Mol, float],
    ):
        metabolite.candidate_structures_eg = candidate_structures

    def set_selected_structure_for_metabolite(
        self,
        metabolite: MappedMetabolite,
        mol: MolTemplate,
        provenance: ProvenanceStruct,
    ):
        metabolite.set_selected_structure(mol, provenance)

    def set_candidate_aams_for_reaction(
        self,
        reaction: MappedReaction,
        candidate_aams: list[tuple[Optional[chemrecon.AAM], ChemReaction, float]],
    ):
        reaction.candidate_aams = candidate_aams

    def add_candidate_matching_for_reaction(
        self, reaction: MappedReaction, candidate_matching: Matching
    ):
        reaction.candidate_matchings.append(candidate_matching)

        # Also add matching information to involved compounds
        for m, mol_template in [
            *candidate_matching.match[Side.L].items(),
            *candidate_matching.match[Side.R].items(),
        ]:
            # Add to the list of structures for this compound
            if m.candidate_structures_aam is None:
                m.candidate_structures_aam = defaultdict(list)
            m.candidate_structures_aam[mol_template].append(
                (reaction, candidate_matching)
            )

    def set_selected_matching(
        self,
        reaction: MappedReaction,
        candidate_matching: Matching,
        provenance: ProvenanceAAM,
    ):
        reaction.selected_aam_matching = candidate_matching
        reaction.selection_provenance = provenance

        for m in reaction.get_metabolites():
            assert isinstance(m, MappedMetabolite)
            if m.candidate_structures_aam_selected is None:
                m.candidate_structures_aam_selected = defaultdict(list)

            struct = candidate_matching.get_structure(m)
            if struct is None:
                # No structure for this metabolite in the selected matching.
                # TODO warn
                continue

            m.candidate_structures_aam_selected[struct].append(
                (reaction, candidate_matching)
            )

    # Problems
    # ----------------------------------------------------------------------------------------------------------
    def add_problem(self, problem: Problem):
        """Add a 'problem' discovered to the ntwork"""
        self.problems.append(problem)
        match problem:
            case MetaboliteProblem(metabolite=m):
                m.problems.append(problem)
            case ReactionProblem(reaction=r):
                r.problems.append(problem)
            case DatabaseAAMProblem(aam_entry=aam):
                # TODO
                pass
            case _:
                # Other
                pass

        print(f" - ! {problem.severity.name}: {problem.message()}")

    # Getters
    # ----------------------------------------------------------------------------------------------------------
    def get_species(self, identifier: str) -> Optional[MappedSpecies]:
        return cast(Optional[MappedSpecies], super().get_species(identifier))

    def get_metabolite(self, identifier: str) -> Optional[MappedMetabolite]:
        return cast(Optional[MappedMetabolite], super().get_metabolite(identifier))

    def get_reaction(self, identifier: str) -> Optional[MappedReaction]:
        return cast(Optional[MappedReaction], super().get_reaction(identifier))

    def get_species_by_id(self, identifier: str) -> Optional[set[MappedSpecies]]:
        species = self._species_std_id_lookup.get(standardize_metabolite_id(identifier))
        if species is None:
            return None
        return species

    def get_metabolite_by_id(self, identifier: str) -> Optional[MappedMetabolite]:
        return self._metabolite_std_id_lookup.get(standardize_metabolite_id(identifier))

    def get_reaction_by_id(self, identifier: str) -> Optional[MappedReaction]:
        return self._reaction_std_id_lookup.get(standardize_reaction_id(identifier))

    # pyrefly: ignore [bad-override]
    def all_species(self) -> list[MappedSpecies]:  # ty: ignore[invalid-method-override]
        return self._species

    # pyrefly: ignore [bad-override]
    def all_metabolites(self) -> list[MappedMetabolite]:  # ty: ignore[invalid-method-override]
        return self._metabolites

    # pyrefly: ignore [bad-override]
    def all_reactions(self) -> list[MappedReaction]:  # ty: ignore[invalid-method-override]
        return self._reactions

    def all_reactions_except_trivial_transport(self) -> list[MappedReaction]:
        # TODO implement upstream?
        return [r for r in self.all_reactions() if not r.is_transport_reaction_1to1]

    # Statistics
    # ----------------------------------------------------------------------------------------------------------
    def stats(self) -> dict[str, Any]:
        stats_dict: dict[str, Any] = dict()

        # Basic stats

        # Metabolite stats
        stats_dict["species"] = len(self.all_species())
        stats_dict["metabolites"] = len(self.all_metabolites())
        stats_dict[" - metabolites (curated)"] = len(self.all_metabolites_curated())
        stats_dict["   - metabolites (curated, hydrogen)"] = len(
            [m for m in self.all_metabolites_curated() if m.is_hydrogen]
        )
        stats_dict[" - metabolites (curated)"] = len(self.all_metabolites_curated())
        stats_dict[" - metabolites (structure from AAM)"] = len(
            self.all_metabolites_implicit_aam_structure()
        )
        stats_dict["    - metabolites (structure from AAM, single)"] = len(
            self.all_metabolites_implicit_aam_structure_single()
        )
        stats_dict["    - metabolites (structure from AAM, multiple)"] = len(
            self.all_metabolites_implicit_aam_structure_multiple()
        )

        # TODO also record degree of conflict?

        stats_dict["  - metabolites (structure from entrygraph)"] = len(
            self.all_metabolites_structure_from_eg()
        )
        stats_dict["  - metabolites (no structure)"] = len(
            self.all_metabolites_no_structure()
        )

        # Reaction stats
        stats_dict["- reactions"] = len(self.all_reactions())

        # Basic stats
        stats_dict["    - reactions (non-transport)"] = len(
            self.all_reactions_nontransport()
        )
        stats_dict["    - reactions (1-to-1 transport)"] = len(
            self.all_reactions_1to1_transport()
        )
        stats_dict["    - reactions (other transport)"] = len(
            self.all_reactions_other_transport()
        )

        # On matching
        stats_dict["    - reactions (curated)"] = len(self.all_reactions_curated())
        stats_dict["    - reactions (database AAM)"] = len(
            self.all_reactions_database_aam()
        )
        stats_dict["    - reactions (generated AAM)"] = len(
            self.all_reactions_generated_aam()
        )
        stats_dict["    - reactions (trivial transport AAM)"] = len(
            self.all_reactions_trivial_transport_aam()
        )
        stats_dict["    - reactions (no AAM)"] = len(self.all_reactions_no_aam())
        stats_dict["       - reactions (no AAM, can generate"] = len(
            self.all_reactions_no_aam_can_generate()
        )
        stats_dict["       - reactions (no AAM, trivial transport)"] = len(
            self.all_reactions_trivial_transport_without_map()
        )
        stats_dict["         - reactions (no AAM, can generate, trivial transport)"] = (
            len(self.all_reactions_trivial_transport_can_generate_aam())
        )

        # TODO matching states

        return stats_dict

    # Metabolite stats
    def all_metabolites_curated(self) -> list[MappedMetabolite]:
        return [
            m
            for m in self.all_metabolites()
            if m.selection_provenance == ProvenanceStruct.USER_CURATED
        ]

    def all_metabolites_implicit_aam_structure(self) -> list[MappedMetabolite]:
        return [
            m
            for m in self.all_metabolites()
            if m.selection_provenance == ProvenanceStruct.IMPLICIT_FROM_AAM
        ]

    def all_metabolites_implicit_aam_structure_single(self) -> list[MappedMetabolite]:
        return [
            m
            for m in self.all_metabolites_implicit_aam_structure()
            if m.candidate_structures_aam is not None
            and len(m.candidate_structures_aam) == 1
        ]

    def all_metabolites_implicit_aam_structure_multiple(self) -> list[MappedMetabolite]:
        return [
            m
            for m in self.all_metabolites_implicit_aam_structure()
            if m.candidate_structures_aam is not None
            and len(m.candidate_structures_aam) > 1
        ]

    def all_metabolites_structure_from_eg(self) -> list[MappedMetabolite]:
        return [
            m
            for m in self.all_metabolites()
            if m.selection_provenance == ProvenanceStruct.VIA_ENTRYGRAPH
        ]

    def all_metabolites_no_structure(self) -> list[MappedMetabolite]:
        return [m for m in self.all_metabolites() if m.selected_mol is None]

    # Reaction stats

    def all_reactions_nontransport(self) -> list[MappedReaction]:
        return [r for r in self.all_reactions() if not r.is_transport_reaction]

    def all_reactions_1to1_transport(self) -> list[MappedReaction]:
        return [r for r in self.all_reactions() if r.is_transport_reaction_1to1]

    def all_reactions_other_transport(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.is_transport_reaction and not r.is_transport_reaction_1to1
        ]

    # Reactions by AAM state
    def all_reactions_curated(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.selection_provenance == ProvenanceAAM.USER_CURATED
        ]

    def all_reactions_database_aam(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.selection_provenance == ProvenanceAAM.IN_DB
        ]

    def all_reactions_generated_aam(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.selection_provenance == ProvenanceAAM.GENERATED
        ]

    def all_reactions_no_aam(self) -> list[MappedReaction]:
        return [r for r in self.all_reactions() if r.selected_aam_matching is None]

    def all_reactions_no_aam_can_generate(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions_no_aam()
            if r.can_generate_reaction and r.selected_aam_matching is None
        ]

    def all_reactions_trivial_transport_aam(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.selection_provenance == ProvenanceAAM.TRIVIAL_TRANSPORT
        ]

    def all_reactions_trivial_transport_without_map(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions()
            if r.is_transport_reaction_1to1 and r.selected_aam_matching is None
        ]

    def all_reactions_trivial_transport_can_generate_aam(self) -> list[MappedReaction]:
        return [
            r
            for r in self.all_reactions_trivial_transport_without_map()
            if r.can_generate_reaction and r.selected_aam_matching is None
        ]

    # Initialization from MetabolicNetwork
    # ----------------------------------------------------------------------------------------------------------
    def __init__(self, network: mntools.MetabolicNetwork):
        """Initialize a mapped network from an existing unmapped network."""

        super().__init__(network.name, network.filename)

        # Create std id lookups
        self._species_std_id_lookup = defaultdict(set)
        self._metabolite_std_id_lookup = dict()
        self._reaction_std_id_lookup = dict()

        self.problems = []

        self._import_from_network(network)

        # Standardize ids
        for ms in self._species:
            self._species_std_id_lookup[standardize_metabolite_id(ms.id)].add(ms)
        for mm in self._metabolites:
            self._metabolite_std_id_lookup[standardize_metabolite_id(mm.id)] = mm
        for mr in self._reactions:
            self._reaction_std_id_lookup[standardize_reaction_id(mr.id)] = mr

    def process(self, unify_metabolites_method: Optional[str] = None):
        """No-op: data was already imported from a pre-processed network in __init__."""
        pass

    def _import_from_network(self, network: mntools.MetabolicNetwork) -> None:
        """Import all entities from a pre-processed MetabolicNetwork as Mapped* equivalents."""
        # Phase 0: network-level fields
        self._compartments = list(network._compartments)
        self._compartment_lookup = dict(network._compartment_lookup)
        self._metabolite_annotation_type = network._metabolite_annotation_type
        self._reaction_annotation_type = network._reaction_annotation_type
        self.warnings = list(network.warnings)

        # Phase 1: MappedSpecies
        species_map: dict[mntools.Species, MappedSpecies] = {}
        for s in network.all_species():
            ms = object.__new__(MappedSpecies)
            ms.__dict__.update(s.__dict__)
            # associated_metabolite and _in_reactions still point to old objects; fixed in phase 4
            species_map[s] = ms
            self._species.append(ms)
            self._species_lookup[ms.id] = ms

        # Phase 2: MappedMetabolite
        metabolite_map: dict[mntools.Metabolite, MappedMetabolite] = {}
        for m in network.all_metabolites():
            mm = MappedMetabolite(
                species=m.species(), unified_id=m.unified_by, unified_by=m.unified_by
            )
            mm.__dict__.update(m.__dict__)
            mm._species_by_compartment = {
                comp: species_map[sp] for comp, sp in m._species_by_compartment.items()
            }
            mm.curated_mol = None
            mm.selected_mol = None
            mm.selection_provenance = None
            mm.selection_details = ""
            mm.candidate_structures_eg = None
            mm.candidate_structures_aam = None
            mm.candidate_structures_aam_selected = None
            mm.is_hydrogen = False
            mm.problems = []
            metabolite_map[m] = mm
            self._metabolites.append(mm)
            self._metabolite_lookup[mm.id] = mm

        # Phase 3: MappedReaction
        reaction_map: dict[mntools.Reaction, MappedReaction] = {}
        for r in network.all_reactions():
            mr = MappedReaction(id=r.id, name=r.name, lhs={}, rhs={})
            mr.__dict__.update(r.__dict__)
            mr._species = {
                side: {species_map[sp]: stoich for sp, stoich in side_dict.items()}
                for side, side_dict in r._species.items()
            }
            mr._metabolites = {
                side: {metabolite_map[met] for met in met_set}
                for side, met_set in r._metabolites.items()
            }
            mr.curated_aam = None
            mr.curated_aam_matching = None
            mr.selected_aam_matching = None
            mr.selection_provenance = None
            mr.selection_details = ""
            mr.candidate_aams = []
            mr.candidate_matchings = []
            reaction_map[r] = mr
            self._reactions.append(mr)
            self._reaction_lookup[mr.id] = mr

        # Phase 4: fix MappedSpecies cross-references
        for s, ms in species_map.items():
            if s.associated_metabolite is not None:
                ms.associated_metabolite = metabolite_map[s.associated_metabolite]
            ms._in_reactions = defaultdict(
                int,
                {reaction_map[r]: stoich for r, stoich in s._in_reactions.items()},
            )
