from __future__ import annotations

from collections import OrderedDict, defaultdict
from typing import TYPE_CHECKING, Callable

import networkx as nx
from chemrecon.chem.mol import Mol, MolTemplate

from aam_annotator.problem import (
    CannotMatch,
    CannotMatchNoMetabolites,
    CannotMatchNoStructures,
)

if TYPE_CHECKING:
    from aam_annotator.mapped_network import MappedMetabolite


def matching_alg(
    metabolites: list[MappedMetabolite],
    structures: list[MolTemplate],
    consider_n_most_confident_strucutres: int = 3,
) -> tuple[list[tuple[MappedMetabolite, MolTemplate]], float]:
    """Main matching algorithm. Will auto-match unmatched metabolites and structures based on information
    in the network.
    """

    if len(metabolites) == 0:
        raise CannotMatchNoMetabolites()
    if len(structures) == 0:
        raise CannotMatchNoStructures()

    metabolite_multiplicity: dict[MappedMetabolite, int] = defaultdict(int)
    structure_multiplicity: dict[MolTemplate, int] = defaultdict(int)
    for m in metabolites:
        metabolite_multiplicity[m] += 1
    for s in structures:
        structure_multiplicity[s] += 1

    # Create graph based on multiplicity
    g = nx.Graph()
    for m, _ in metabolite_multiplicity.items():
        g.add_node(m)
    for struct, _ in structure_multiplicity.items():
        g.add_node(struct)

    # Add edges based on compatibility
    for i, m in enumerate(metabolites):
        for struct in structures:
            # Add edge with compatibility as edge
            try:
                comp = compatibility(
                    m,
                    struct,
                    compatibility_measure=compatibility_by_conf_weighted_similiarity,
                    consider_n_most_confident_structures=consider_n_most_confident_strucutres,
                )
            except ValueError:
                # No candidate structures found
                comp = 0

            # Modify depending on multiplicity
            if metabolite_multiplicity[m] != structure_multiplicity[struct]:
                # Multiplicity mismatch - compatibility is lowered
                # TODO details?
                comp *= 0.2

            g.add_edge((i, m), struct, weight=comp)

    # Do checks
    # - if all edges have weight 0, return empty list
    if all(w == 0 for _, _, w in g.edges(data="weight")):
        # No structures to match.
        raise CannotMatch()

    # Run graph matching alg
    graph_matching = nx.max_weight_matching(
        g, maxcardinality=False
    )  # TODO experiment with max cardinality.

    # Interpret result (returned as a list not necessarily ordered)
    matching: list[tuple[MappedMetabolite, MolTemplate]] = list()
    keys: set = set(enumerate(metabolites))
    for a, b in graph_matching:
        if a in keys:
            matching.append((a[1], b))
        elif b in keys:
            matching.append((b[1], a))
        else:
            raise RuntimeError("Graph matching returned an invalid result!")

    # Compute confidence of match and return
    confsum = sum(g.get_edge_data(a, b)["weight"] for a, b in graph_matching)
    conf = confsum / len(graph_matching) if len(matching) != 0 else 1
    return matching, conf


# Compatibility
def compatibility(
    metabolite: MappedMetabolite,
    mol: Mol,
    compatibility_measure: Callable[[dict[Mol, float], dict[Mol, float], int], float],
    consider_n_most_confident_structures: int = 3,
) -> float:
    """Compute the degree of database-based compatibility between a network metabolite and a molecular
    structure. Essentially: How likely is this match?
    """
    # TODO weight primary higher in scoring
    # Find score of each molecular structure in the EGs, converting db structs to internal repr.
    scores: OrderedDict[Mol, float]
    if metabolite.curated_mol is None:
        # Get compounds by EntryGraph
        cands = metabolite.candidate_structures_eg
        if cands is None:
            raise ValueError("No candidate structures found!")
        scores = cands
    else:
        # Curated molecular structure given, consider only this
        scores = OrderedDict({metabolite.curated_mol: 1.0})

    # Compute similarities
    if len(scores) == 0:
        # No structures found
        return 0

    return compatibility_measure(
        scores,
        {
            mol_: mol.get_similarity(mol_)
            for mol_ in list(scores.keys())[:consider_n_most_confident_structures]
        },
        consider_n_most_confident_structures,
    )


def similarity(m: Mol, m_: Mol) -> float:
    return m.get_similarity(m_)


# Compatibility measure functions
# ----------------------------------------------------------------------------------------------------------------------
# We can assume that len(eg_scores), len(similarity) > 0
def compatibility_by_conf_weighted_similiarity(
    eg_scores: dict[Mol, float],
    sim: dict[Mol, float],
    take_first_n: int,
) -> float:
    # Average similarity, weighted by eg score
    return sum(
        sim[m] * conf for m, conf in list(eg_scores.items())[:take_first_n]
    ) / sum(list(eg_scores.values())[:take_first_n])


def compatibility_max_by_eg_conf(
    eg_scores: dict[Mol, float],
    sim: dict[Mol, float],
    take_first_n: int,
) -> float:
    sims: list[float] = list()
    for m, conf in list(eg_scores.items())[:take_first_n]:
        sims.append(sim[m] * conf)
    return max(sims) if (len(sims) > 0) else 0
