"""Main algorithm."""

from collections import defaultdict
from pathlib import Path
from typing import Optional

import chemrecon
import mntools
from chemrecon import ChemReaction
from chemrecon.chem.chemreaction import (
    chem_reaction_from_aam_entry,
    chem_reaction_from_aam_entry_flipped,
)
from chemrecon.chem.mol import Mol, MolTemplate

from aam_annotator.atom_to_atom_mapper import produce_aam
from aam_annotator.curations import Curations, load_curations_from_directory
from aam_annotator.mapped_network import MappedMetabolite, MappedNetwork
from aam_annotator.matching import (
    DatabaseAAMMatching,
    GeneratedMatching,
    Matching,
    TrivialTransportMatching,
    UserCuratedAAMMatching,
)
from aam_annotator.problem import (
    AAMUnbalancedProblem,
    AnotherMatchingHasHigherConfidence,
    CannotGenerateAAMNoStructures,
    CannotMatch,
    CouldNotGenerateMapToCanonicalStructure,
    CuratedAAMNotFound,
    CuratedStructureNotFound,
    IncompatibleStructuresSelectedWarning,
    IncompleteMatchingProblem,
    LowConfidenceMatchingWarning,
    MissingStructuresWarning,
    NoAAMsFound,
    NoStructuresFound,
    RDTError,
)
from aam_annotator.provenance import ProvenanceAAM, ProvenanceStruct
from aam_annotator.search import get_eg_aams, get_eg_structures
from aam_annotator.serialize import serialize_network
from aam_annotator.utils.std_id import standardize_metabolite_id

Side = mntools.Side


def generate_mapped_network(
    network: mntools.MetabolicNetwork,
    output_filename: Optional[str] = None,
    aam_search_depth: int = 4,
    structure_search_depth: int = 3,
    generate_aams: bool = True,
    curations_path: Optional[str] = "./curations/",
) -> MappedNetwork:

    # First, test whether ChemRecon is connected. Otherwise, connect to the public database server.
    if chemrecon.get_query_handler() is None:
        chemrecon.connect_public()

    # Initialize the mapped network
    # ==========================================================================================================
    print("\nInitializing the mapped network...")
    mapped_network = MappedNetwork(network)
    mapped_network.process()

    # Load curations
    # ==========================================================================================================
    print("\nLoading curations...")
    if curations_path is not None:
        curations = load_curations_from_directory(curations_path)
    else:
        curations = Curations()

    # Metabolites (structures)
    for std_id, mol in curations.structures.items():
        m = mapped_network.get_metabolite_by_id(std_id)
        if m is not None:
            m.set_curated_structure(mol)
            curations.structure_matchings[std_id] = m

            # Hydrogen
            if std_id in curations.hydrogen_metabolites:
                m.is_hydrogen = True

        else:
            print(f" - No metabolite found for {std_id}!")
            mapped_network.add_problem(CuratedStructureNotFound(std_id, mol))
            continue

    # Reactions (AAMs)
    for std_id, aam in curations.aams.items():
        r = mapped_network.get_reaction_by_id(std_id)
        if r is not None:
            r.set_curated_aam(aam)
            curations.reaction_matchings[std_id] = r
        else:
            print(f" - No reaction found for {std_id}!")
            mapped_network.add_problem(CuratedAAMNotFound(std_id, aam))
            continue

    # Find structures for each metabolite
    # ==========================================================================================================
    print(
        f"\nFinding structures for {len(mapped_network.all_metabolites())} metabolites..."
    )
    for i, m in enumerate(mapped_network.all_metabolites()):
        print(f" - {i + 1}/{len(mapped_network.all_metabolites())}: {m.id}")

        if m.curated_mol is not None:
            print("   - Already curated.")
            continue

        if m.is_biomass:
            print("   - Is biomass, skipping.")
            continue

        if m.is_hydrogen:
            print("   - Is hydrogen, skipping.")
            continue

        # Perform EG-based lookup (cached)
        try:
            molstructure_entries = get_eg_structures(
                m,
                search_depth=structure_search_depth,
                cache_key=standardize_metabolite_id(m.id),
            )
            mapped_network.set_candidate_structures_for_metabolite(
                m, molstructure_entries
            )
            print(f"   - Found {len(molstructure_entries)} structures")

        except NoStructuresFound as p:
            mapped_network.add_problem(p)
            continue

    # Find AAMs for each reaction
    # ==========================================================================================================
    print(f"\nFinding AAMs for {len(mapped_network.all_reactions())} reactions...")
    for i, r in enumerate(mapped_network.all_reactions()):
        print(f" - {i + 1}/{len(mapped_network.all_reactions())}: {r.id}")

        if r.curated_aam is not None:
            print("   - Already curated.")
            continue

        if r.is_biomass_reaction:
            print("   - Is biomass reaction, skipping.")
            continue

        # Perform EG-based lookup (cached)
        try:
            aam_entries = get_eg_aams(r, search_depth=aam_search_depth, cache_key=r.id)
            print(f"   - Found {len(aam_entries)} AAMs")
            chem_reactions: list[
                tuple[Optional[chemrecon.AAM], ChemReaction, float]
            ] = []
            for aam, score in aam_entries.items():
                # Add both flipped and unflipped versions to candidate reactions.
                chem_reactions.append((aam, chem_reaction_from_aam_entry(aam), score))
                chem_reactions.append(
                    (aam, chem_reaction_from_aam_entry_flipped(aam), score)
                )

            # Add as candidates
            mapped_network.set_candidate_aams_for_reaction(r, chem_reactions)

        except NoAAMsFound as p:
            mapped_network.add_problem(p)
            continue

    # Produce candidate matchings for the candidate AAMs.
    # ==========================================================================================================
    print("\nGenerating candidate matchings for each AAM...")
    for i, r in enumerate(mapped_network.all_reactions_except_trivial_transport()):
        print(f" - {i + 1}{r.id}")

        # Get compound entries for each metabolite linked in the CRN reaction
        compound_entries: dict[
            Side, dict[MappedMetabolite, set[chemrecon.Compound]]
        ] = {Side.L: dict(), Side.R: dict()}

        for m, _ in r.get_lhs_metabolites().items():
            assert isinstance(m, MappedMetabolite)
            compound_entries[Side.L][m] = m.entries
        for m, _ in r.get_rhs_metabolites().items():
            assert isinstance(m, MappedMetabolite)
            compound_entries[Side.R][m] = m.entries

        # If curated AAM
        if r.curated_aam is not None:
            # If curated, set AAM entry to NONE (will be used in next step)
            r.candidate_aams = [(None, r.curated_aam, 1.0)]

        # For each candidate AAM, generate a candidate matching.
        for aam_index, (candidate_aam_entry, candidate_chemreaction, conf) in enumerate(
            r.candidate_aams
        ):
            print(f"   - AAM candidate {aam_index + 1}/{len(r.candidate_aams)}")

            if not candidate_chemreaction.is_balanced_up_to_h():
                mapped_network.add_problem(
                    AAMUnbalancedProblem(r, candidate_aam_entry, candidate_chemreaction)
                )

            # Create matchings
            if candidate_aam_entry is not None:
                # AAM from DB
                matching = DatabaseAAMMatching(
                    r, candidate_chemreaction, candidate_aam_entry, conf
                )
            else:
                # AAM from curation
                matching = UserCuratedAAMMatching(r, candidate_chemreaction)

            # Apply curations and match the rest
            matching.match_from_curations()
            if matching.is_complete():
                # Already complete
                pass
            else:
                # Else, match the rest by heuristic
                try:
                    matching.match_structures()
                except CannotMatch as p:
                    matching.add_problem(p)

            # Note if incomplete
            if not matching.is_complete():
                matching.add_problem(IncompleteMatchingProblem())

            # Add as candidate matching
            matching.process()
            mapped_network.add_candidate_matching_for_reaction(r, matching)

    # Choose an AAM and a matching for each reaction
    # ==========================================================================================================
    print("\nChoosing an AAM and a matching for each reaction...")
    # Sort by best matching conf
    sorted_matched_reactions = sorted(
        [
            r
            for r in mapped_network.all_reactions_except_trivial_transport()
            if len(r.candidate_matchings) > 0
        ],
        key=lambda r_: max(m_.get_conf() for m_ in r_.candidate_matchings),
        reverse=True,  # Descending
    )
    for r in sorted_matched_reactions:
        print(f" - {r.id}")
        sorted_matchings = sorted(
            r.candidate_matchings, key=lambda m_: m_.get_conf(), reverse=True
        )

        # Recognize metabolite-structure conflicts with neighbouring reactions
        match_neighbour_conflicts: dict[Matching, set[MappedMetabolite]] = dict()
        for matching in r.candidate_matchings:
            conflicting_metabolites: set[MappedMetabolite] = set()
            for m in r.get_metabolites():
                assert isinstance(m, MappedMetabolite)
                mol = matching.get_structure(m)

                # Check that this is assigned
                if m.candidate_structures_aam_selected is None:
                    # No assigned AAM structures yet, so no conflict
                    pass
                else:
                    # Add conflict if matching struct not already in assigned set.
                    assigned_structs = m.candidate_structures_aam_selected.keys()
                    if m.selected_mol in assigned_structs:
                        conflicting_metabolites.add(m)

            if len(conflicting_metabolites) > 0:
                match_neighbour_conflicts[matching] = conflicting_metabolites

        # Consider matchings in ascending order of neighbour conflicts, then by descending confidence
        candidate_matchings_sorted: list[Matching] = sorted(
            sorted(
                # Second, sort by confidence
                r.candidate_matchings,
                key=lambda m_: m_.get_conf(),
                reverse=True,
            ),
            # First, sort by number of conflicts
            key=lambda m_: len(match_neighbour_conflicts.get(m_, set())),
        )

        # Assign best AAM matching
        if len(candidate_matchings_sorted) == 0:
            print(f"   - No matchings found for {r.id}!")
            continue
        selected_matching = candidate_matchings_sorted[0]

        provenance: ProvenanceAAM
        if isinstance(selected_matching, DatabaseAAMMatching):
            provenance = ProvenanceAAM.IN_DB
        elif isinstance(selected_matching, UserCuratedAAMMatching):
            provenance = ProvenanceAAM.USER_CURATED
        else:
            raise NotImplementedError("Should not happen.")

        mapped_network.set_selected_matching(r, selected_matching, provenance)
        selected_conf = selected_matching.get_conf()

        print("   - Selected matching:")
        print(f"     -> {selected_matching}")
        print(f"     -> Confidence: {selected_conf:.2f}")

        # Warn if matching is not confident
        if selected_conf < 0.5:
            selected_matching.add_problem(LowConfidenceMatchingWarning(selected_conf))

        # Warn if mirror matching has similar confidence
        # TODO

        # Warn if another matching has greater conf
        max_conf_matching: Matching | None = None
        for m_ in candidate_matchings_sorted:
            if max_conf_matching is not None:
                if m_.get_conf() > max_conf_matching.get_conf():
                    max_conf_matching = m_
            else:
                max_conf_matching = m_
        if max_conf_matching is not None and max_conf_matching != selected_matching:
            selected_matching.add_problem(
                AnotherMatchingHasHigherConfidence(max_conf_matching)
            )

    # Consolidate structures for each metabolite, and choose a representative structure
    # ==========================================================================================================
    print("\nConsolidating structures for each metabolite...")
    for i, m in enumerate(mapped_network.all_metabolites()):
        print(f"- {i + 1}/{len(mapped_network.all_metabolites())}: {m}")

        if m.curated_mol is not None:
            print("   - Already curated!")
            # TODO investigate whether curated mol matches aams
            continue

        # Find structures from AAMs
        aam_structures_support: dict[MolTemplate, int] = (
            {
                moltemp: len(support)
                for moltemp, support in m.candidate_structures_aam_selected.items()
            }
            if m.candidate_structures_aam_selected is not None
            else {}
        )

        if len(aam_structures_support) == 0:
            print(f"   - No AAM structures found for {m.id}!")
            # No aam structures
            continue

        # Group structures by maximally standardized structure.
        std_dict: dict[Mol, set[MolTemplate]] = defaultdict(set)
        for mol, _ in aam_structures_support.items():
            mol_std = mol.get_standardised()
            std_dict[mol_std].add(mol)

        # Get most supported set
        greatest_std_mol, greatest_support_set = max(
            std_dict.items(),
            key=lambda item: sum(
                aam_structures_support[moltemp] for moltemp in item[1]
            ),
        )

        # Representative mol is the most standardized structure in the greatest support set
        all_features = set().union(*(mol.features for mol in greatest_support_set))
        representative_mols = [
            mol for mol in greatest_support_set if mol.features == all_features
        ]
        representative_mol: Mol
        if len(representative_mols) > 1:
            # There is no maximum in the feature lattice

            # Get the structure most supported by the entry graph
            if m.candidate_structures_eg is not None:
                representative_mol_eg_conf: dict[MolTemplate, float] = dict()
                for rep_mol in representative_mols:
                    try:
                        conf = m.candidate_structures_eg[rep_mol]
                    except KeyError:
                        conf = 0.0
                    representative_mol_eg_conf[rep_mol] = conf

                # Now, the representative mol is the one with the highest confidence
                representative_mol = max(
                    representative_mol_eg_conf,
                    key=lambda k: representative_mol_eg_conf[k],
                )

            else:
                # If we cannot select by EG conf
                # There is no Maximum in the feature lattice, choose the max standardized
                assert greatest_std_mol is not None
                representative_mol = greatest_std_mol

            # Get conflicting structures
            conflicting_mols: list[Mol] = []
            assert m.candidate_structures_aam_selected is not None
            for mol in m.candidate_structures_aam_selected.keys():
                # Standardize to see if consistent with the representative mol
                mol_std = mol.get_standardised(list(representative_mol.features))
                # Mols are over-standardized here!

                if mol_std != representative_mol:
                    conflicting_mols.append(mol)

            if len(conflicting_mols) > 0:
                mapped_network.add_problem(
                    IncompatibleStructuresSelectedWarning(
                        m, representative_mol, conflicting_mols
                    )
                )
            else:
                # Should this happen?
                pass

        elif len(representative_mols) == 1:
            # Choose this
            representative_mol = representative_mols[0]
        else:
            # Should not happen
            raise NotImplementedError("Should not happen.")

        # Assign the representative mol as the selected mol
        print(f"   - Selected representative mol: {representative_mol}")
        representative_mol_template = MolTemplate(representative_mol.mol)
        m.set_selected_structure(
            representative_mol_template, ProvenanceStruct.IMPLICIT_FROM_AAM
        )

        # Produce maps to the representative mol
        # ------------------------------------------------------------------------------------------------------
        if (m.candidate_structures_aam_selected is not None) and (
            len(m.candidate_structures_aam_selected) > 1
        ):
            print(
                f"   - Producing canonical maps for {m.id}: {len(m.candidate_structures_aam_selected)} structures"
            )
            for alt_struct in m.candidate_structures_aam_selected:
                assert isinstance(m.selected_mol, MolTemplate)
                if alt_struct == m.selected_mol:
                    continue

                # Check similarity?
                # TODO

                # Produce one-to-one AAM
                canonical_map_aam = produce_aam(lhs=[alt_struct], rhs=[m.selected_mol])
                if canonical_map_aam is not None:
                    m.set_canonical_map_aam(alt_struct, canonical_map_aam)
                else:
                    # Include error message with query
                    smiles_lhs = alt_struct.to_smiles()
                    smiles_rhs = m.selected_mol.to_smiles()
                    rdt_query: str = smiles_lhs + ">>" + smiles_rhs
                    mapped_network.add_problem(
                        CouldNotGenerateMapToCanonicalStructure(
                            m,
                            rdt_query,
                        )
                    )

    # Assign structures to metabolites with no AAM structures
    # ==========================================================================================================
    print("\nAssigning structures to metabolites with no AAM structures...")
    for m in mapped_network.all_metabolites():
        if m.selected_mol is not None:
            continue
        print(f" - {m}")

        # Get EG structures
        match m.candidate_structures_eg:
            case None:
                # No EG structures found
                # TODO warn
                print("   - No EG structures found")
                continue
            case eg_structures:
                # Use most confident for now
                assert eg_structures is not None
                if len(eg_structures) == 0:
                    continue
                print(f"   - Found {len(eg_structures)} EG structures")
                best_eg_mol = max(eg_structures, key=lambda k: eg_structures[k])

                # Use selected mol
                print(
                    f"   - Using confidence structure: {best_eg_mol} ({eg_structures.get(best_eg_mol):.2f})"
                )
                assert isinstance(best_eg_mol, MolTemplate)
                m.set_selected_structure(best_eg_mol, ProvenanceStruct.VIA_ENTRYGRAPH)

    # Process reactions with missing maps
    # ==========================================================================================================
    print("\nProcessing reactions with missing maps")
    for r in mapped_network.all_reactions():
        print(f" - {r}")

        if r.is_biomass_reaction:
            print("   - Is biomass reaction, skipping.")
            continue

        # Get structures for each compound1
        missing_mols: set[MappedMetabolite] = set()
        mols: dict[Side, list[Mol]] = {Side.L: [], Side.R: []}
        for side, items in [
            (Side.L, r.get_lhs_metabolites().items()),
            (Side.R, r.get_rhs_metabolites().items()),
        ]:
            for m, stoich in items:
                assert isinstance(m, MappedMetabolite)
                if m.selected_mol is None:
                    # Missing
                    missing_mols.add(m)
                else:
                    # Add
                    for _ in range(abs(stoich)):
                        mols[side].append(m.selected_mol)

        if len(missing_mols) > 0:
            # If missing all structures on either side
            if len(mols[Side.L]) == 0 or len(mols[Side.R]) == 0:
                mapped_network.add_problem(
                    CannotGenerateAAMNoStructures(r, list(missing_mols))
                )
                continue
            else:
                # If not critical, just emit a warning
                # TODO option to skip generating AAMs in this case?
                mapped_network.add_problem(
                    MissingStructuresWarning(r, list(missing_mols))
                )

        # Produce trivial matching based on smiles
        r.can_generate_reaction = True
        if generate_aams:
            aam = produce_aam(mols[Side.L], mols[Side.R])
            if aam is None:
                # Could not generate map, compute the query and add a problem
                smiles_lhs = ".".join(
                    sorted([temp.to_smiles() for temp in mols[Side.L]])
                )
                smiles_rhs = ".".join(
                    sorted([temp.to_smiles() for temp in mols[Side.R]])
                )
                rdt_query: str = smiles_lhs + ">>" + smiles_rhs
                mapped_network.add_problem(RDTError(r, rdt_query))
                continue

            # Produce matching
            if r.is_transport_reaction_1to1:
                # If 1-to-1 transport reaction, use trivial matching
                matching = TrivialTransportMatching(r, aam)
            else:
                # Else, match based on the generated AAM
                matching = GeneratedMatching(r, aam)

            # Add matching, and set as selected
            matching.match_structures()
            matching.process()
            mapped_network.add_candidate_matching_for_reaction(r, matching)
            mapped_network.set_selected_matching(r, matching, ProvenanceAAM.GENERATED)
        else:
            print("   -> AAM generation off, skipping.")

    # Analyze, collect potential problems, ...
    # ==========================================================================================================

    # Identify matchings mapping to non-primary structures
    # TODO add problems for matchings that map to non-primary structures

    # Serialize and output
    # ==========================================================================================================
    # Output stats
    print("\nDone.")
    print("\nStatistics:")
    for key, value in mapped_network.stats().items():
        print(f"  {key}: {value}")

    json_str = serialize_network(mapped_network, curations)
    if output_filename is not None:
        output_path = Path(output_filename)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(json_str)

    # Done.
    # ==========================================================================================================

    return mapped_network
