from __future__ import annotations

import networkx as nx

from chem.elements import ChemGraph, chemgraph_from_mol, ematch, nmatch
from chem.elements import mol_from_structurerepresentation
from chemrecon.core.query_handler import QueryHandler
from chemrecon.schema.entry_types.compound import Compound
from chemrecon.schema.entry_types.reaction import Reaction
from chemrecon.schema.relation_types_source.compound_has_structure_representation_relation import \
    CompoundHasStructureRepresentation
from chemrecon.schema.relation_types_source.reaction_involves_compound_relation import ReactionInvolvesCompound
from utils import hungarian as hungarian


# TODO change to use the general reactant_matching with ChemReactions

def match_reactants(
    handler: QueryHandler,
    procedure: Procedure,
    chemgraph: ChemGraph,
    reaction: Reaction,
    side: int  # -1 or 1
) -> dict[frozenset[int], tuple[Compound, int]]:
    """ Given a left or right S_GML of a reaction (already sanitised by removing enzymes), tries to match the
        components of the graph to a list of entries with known structures using graph similarity. Returns a
        mapping, from each reactant (compound, multiplicity index), to the set of nodes with which they
        were identified.
    """

    # TODO eventually replace the structure getting alg. with entry graph-based approach

    # Get participant compounds
    reaction_all_compounds: list[tuple[Compound, int]] = list()  # With stoich
    compound_result_dict = handler.get_relations_with_entries_by_recon_ids_of_t1(
        entry_type = Reaction,
        recon_ids = [reaction.recon_id],
        relation_type = ReactionInvolvesCompound,
    )[0]
    for compound_entry, relations in compound_result_dict.items():
        relations: list[ReactionInvolvesCompound]
        for relation in relations:
            reaction_all_compounds.append(
                (compound_entry, relation.n)
            )

    # Get existing structures for the attached compounds (converted to chemgraph objects?)
    compound_graphs: dict[Compound, ChemGraph] = dict()
    for compound_entry, stoich in reaction_all_compounds:
        # Only take compounds on the correct side
        if stoich * side < 0:
            continue

        # Get structures
        struct_repr_relations = handler.get_relations_with_entries_by_recon_ids_of_t1(
            entry_type = Compound,
            recon_ids = [compound_entry.recon_id],
            relation_type = CompoundHasStructureRepresentation
        )
        struct_reprs = struct_repr_relations[0].keys()

        # For now, select an arbitrary structure
        try:
            struct_repr = list(struct_reprs)[0]
        except IndexError:
            # No structures found
            continue

        # Convert the selected structure to a S_GML ChemGraph
        compound_graphs[compound_entry] = chemgraph_from_mol(mol_from_structurerepresentation(struct_repr))

    # Perform the matching algorithm on match_structures against components of the given chemgraph
    chemgraph.remove_hydrogens()
    component_subsets: list[frozenset[int]] = chemgraph.get_component_subsets()
    component_subgraphs: list[nx.Graph] = chemgraph.get_components_subgraphs()

    # Log an error if component numbers don't match
    if len(component_subgraphs) != len(compound_graphs):
        # TODO Try to fix this by balancing cofactors etc
        # TODO May be extra hydrons, water etc
        # TODO  - entry #6 has an extra water on both sides for instance
        procedure.log_notice(f'Could not get structures for all components.'
                             f' Reaction {reaction} (side {side}).'
                             f' Number of S_GML graph components ({len(component_subgraphs)})'
                             f' does not match number of reactants ({len(compound_graphs)}).')

    # Pairwise similarity (maps from each compound to the index of the component and the similarity)
    graph_distance_pairwise: dict[Compound, list[tuple[int, float]]] = dict()

    # Compute the similarity
    for compound_entry, compound_chemgraph in compound_graphs.items():
        component_distances: list[tuple[int, float]] = list()
        for component_index, component in enumerate(component_subgraphs):
            # nx.optimize_graph_edit_distance is a generator which generates successively more precise values.
            dist = next(nx.optimize_graph_edit_distance(
                compound_chemgraph.graph,
                component,
                node_match = nmatch,
                edge_match = ematch
            ))
            component_distances.append((component_index, dist))
        graph_distance_pairwise[compound_entry] = component_distances

    # We make duplicate instances of each compound to match if stoich != 1
    compound_instances: dict[tuple[Compound, int], ChemGraph] = dict()
    for compound, stoich in reaction_all_compounds:
        if stoich * side < 0:
            continue
        for i in range(abs(stoich)):
            try:
                compound_instances[(compound, i)] = compound_graphs[compound]
            except KeyError as e:
                pass

    # Run the hungarian algorithm to get the best match with the instanced compounds
    edges: dict[tuple[tuple[Compound, int], int], float] = dict()
    for (compound, compound_instance_index), _ in compound_instances.items():
        for component_index, distance in graph_distance_pairwise[compound]:
            edges[((compound, compound_instance_index), component_index)] = distance
    matching = hungarian.max_weight_matching(
        edges = edges,
        min_weight = True
    )
    # TODO extra logic to iterate the edit distance algorithm if the matching is 'close'?

    matching_dict: dict[set[int], tuple[Compound, int]] = dict()
    for (compound, compound_index), component_index in matching.items():
        matching_dict[component_subsets[component_index]] = (compound, compound_index)

    return matching_dict
