""" Implements a heuristic algorithm to match structures of a ChemReaction to the structures attached to a
    Compound.
"""
from __future__ import annotations

from collections import OrderedDict
from enum import Enum
from typing import Hashable

from chemrecon import Direction, ChemReaction, Mol
from chemrecon.chem.create_mol import mol_from_struct_entry
from chemrecon.entrygraph.entrygraph import EntryGraph
from chemrecon.entrygraph.explorationprotocol import ExplorationProtocol
from chemrecon.entrygraph.explore import explore
from chemrecon.entrygraph.scoring import Scorer
from chemrecon.schema import (
    Compound, MolStructure,
)
from chemrecon.schema.relation_types_procedural.compound_select_structure_proceduralrelation import \
    CompoundSelectStructure
from chemrecon.utils import hungarian


class FlipState(Enum):
    NORMAL = 0
    FLIPPED = 1


def match_reactants(
        reaction: ChemReaction,
        compound_entries_lhs: dict[str, list[Compound]],
        compound_entries_rhs: dict[str, list[Compound]],
        consider_n_most_confident_structures: int = 1,
        consider_first_entry_primary: bool = True,
        consider_flipped: bool = False,
        search_depth: int = 2
) -> tuple[dict[Mol, str], dict[Mol, str], FlipState]:
    """ ...
        Compound entries are given as a dict (with key = primary_id), and a list of associated entries, possibly only 1.
        If consider_first_entry_primary is set, for each given compound as a list of entries, the first entry in the
        list will have higher weight in the EntryGraph.
        Each compound can be given with multiple entries (for SBML files, for instance).
        If consider_flipped is set, will also consider the case when LHS and RHS are flipped.
        If look_deep is set, will consider structures not directly associated with the entries, instead using a deeper
        Entrygraph search.
        Also returns whether flipped was best
    """

    # First, find structures using EntryGraphs
    compound_structures_lhs: dict[str, OrderedDict[Mol, float]] = OrderedDict()
    compound_structures_rhs: dict[str, OrderedDict[Mol, float]] = OrderedDict()
    for compound_entries, compound_structures in [
        [compound_entries_lhs, compound_structures_lhs],
        [compound_entries_rhs, compound_structures_rhs]
    ]:
        compound_entries: dict[str, list[Compound]]
        compound_structures: dict[str, OrderedDict[Mol, float]]

        for k, entry_list in compound_entries.items():
            try:
                compound_structures[k] = dict()
                for struct_entry, score in search_structures(entry_list).items():
                    try:
                        compound_structures[k][mol_from_struct_entry(struct_entry)] = score
                    except AttributeError:
                        # Could not generate Mol
                        pass
            except ValueError:
                # No valid initial entries
                compound_structures[k] = dict()

    # Perform matching based on the Mols in the ChemReaction
    match_l: dict[Mol, str]
    match_r: dict[Mol, str]
    match_ll, score_ll = match(
        reaction.get_lhs_templates(),
        compound_structures_lhs,
        consider_n_most_confident_structures = consider_n_most_confident_structures
    )
    match_rr, score_rr = match(
        reaction.get_rhs_templates(),
        compound_structures_rhs,
        consider_n_most_confident_structures = consider_n_most_confident_structures
    )

    # Consider the reaction flipped
    if consider_flipped:
        match_lr, score_lr = match(
            reaction.get_lhs_templates(),
            compound_structures_rhs,
            consider_n_most_confident_structures = consider_n_most_confident_structures
        )
        match_rl, score_rl = match(
            reaction.get_rhs_templates(),
            compound_structures_lhs,
            consider_n_most_confident_structures = consider_n_most_confident_structures
        )

        if (score_lr + score_rl) > (score_ll + score_rr):
            print(f'  -> flipped best ({score_lr:.2f} + {score_rl:.2f}) > ({score_ll:.2f} + {score_rr:.2f})')
            flipstate = FlipState.FLIPPED
            match_l = match_lr
            match_r = match_rl
        else:
            print(f'  -> non-flipped best ({score_lr:.2f} + {score_rl:.2f}) < ({score_ll:.2f} + {score_rr:.2f})')
            flipstate = FlipState.NORMAL
            match_l = match_ll
            match_r = match_rr
    else:
        flipstate = FlipState.NORMAL
        match_l = match_ll
        match_r = match_rr

    # Return chosen matching
    return (match_l, match_r, flipstate)


# Matching alg
def match[T: Hashable](
        reaction_mols: list[Mol],
        compound_mols: dict[T, OrderedDict[Mol, float]],
        consider_n_most_confident_structures: int = 1,
) -> tuple[dict[Mol, T], int]:
    """ Given the mols from the reaction and the compound_mols (with confidence), get the best mapping.
        Type parameter T is the key, probalby str (primary_id).
        Returns the matching, as well as a confidence level (0 to 1).
    """
    similarity: dict[tuple[Mol, T], float] = dict()

    # First, compute the pairwise similarity of all compounds
    for mol in reaction_mols:
        for primary_id, mols_ in compound_mols.items():
            # if len(mols_) == 0:
            #    # No structures given
            #    continue

            sims: list[float] = list()
            for mol_, conf in list(mols_.items())[:consider_n_most_confident_structures]:
                # Get distance
                sims.append(mol.get_similarity(mol_) * conf)  # 1 if identical, approaches 0 as difference increases.

            # Calculate the similarity as the maximum of (sim * conf)
            similarity[(mol, primary_id)] = max(sims) if sims else 0

    # Run the Hungarian algorithm to determine the best matching
    matching = hungarian.max_weight_matching(edges = similarity)

    # Compute the confidence/score of the matching as avg. weight of matched edges
    confidence = sum(similarity[mol, primary_id] for mol, primary_id in matching.items()) / len(reaction_mols)
    return matching, confidence


def search_structures(compound_entries: list[Compound]) -> OrderedDict[MolStructure, float]:
    eg = EntryGraph(initial_entries = set(compound_entries))
    explore(eg, structure_protocol, steps = 4)
    return scorer_structure(eg)


# EntryGraph specifications for search
# ----------------------------------------------------------------------------------------------------------------------
structure_protocol = ExplorationProtocol(
    relation_types = {CompoundSelectStructure}
)
scorer_structure = Scorer[MolStructure](
    score_entry_type = MolStructure
)

# TODO score depending on whether primary id is weighted higher
