""" Methods for producing a ranking on an entry graph.
"""
from collections import OrderedDict
from typing import Callable, Optional, TYPE_CHECKING

import rustworkx

from chemrecon.entrygraph.entrygraph import (
    Edge, EntryGraph, SourceEdgeArtificial, SourceVertexArtificial, Vertex,
)
from chemrecon.schema import Entry, Relation

if TYPE_CHECKING:
    from chemrecon.schema.procedural_relation_entrygraph import ProceduralRelationEG


class Scorer[T_rank: Entry]:
    """ A scorer is a callable which takes an entrygraph and produces a ranking of the vertices according to
        the parameters of the scorer.

        The score of an entry is (informally) the probability that a random walk starting at one of the initial entries
        of the entry graph will terminate at that entry.
        The parameters of the random walk can be customized by specifying weights (probabilities) using a weight
        function on entries and relations, which alters the probability of choosing a given path.
        The default weight of all entries and relations is 1.
        For example, if you do not trust a particular source, edges and vertices from that source can have their
        weight reduced, making them count less in the scoring algorithm.

        A damping factor, `alpha` can be specified.
        With probability `1-alpha`, the random walk will choose to go to a random entry rather than continuing the walk.
        Furthermore, a _decay_factor_ can be specified such that entries further away from the initial entries are given
        lower scores.
        A decay factor of `0` disables this adjustment. The default is `0.2`.

        Scores are normalized such that the sum of scores is **1**, which allows comparing scores across entry graphs.

        Formally, the scores are computed using the PageRank algorithm (https://en.wikipedia.org/wiki/PageRank),
        starting from the initial vertices, and with dangling vertices pointing back to all initial vertices with
        equal probability.

    """
    score_entry_type: type[T_rank]

    entry_weight: Callable[[Entry], float]
    relation_weight: Callable[[Relation], float]

    # Algorithm parameters
    alpha: float
    decay_factor: float

    def __init__(
            self,
            score_entry_type: type[Entry],
            alpha: float = 0.85,
            decay_factor: float = 0.2,
            entry_weight: Optional[Callable[[Entry], float]] = None,
            relation_weight: Optional[Callable[[Relation], float]] = None,
    ):
        """ Specify a scorer.
        """
        self.alpha = alpha
        self.decay_factor = decay_factor

        self.score_entry_type = score_entry_type

        if entry_weight is not None:
            self.entry_weight = entry_weight
        else:
            self.entry_weight = lambda e: 1

        if relation_weight is not None:
            self.relation_weight = relation_weight
        else:
            self.relation_weight = lambda r: 1

    def __call__(self, entrygraph: EntryGraph) -> OrderedDict[T_rank, float]:
        """ Produces a ranking of the entries of the type `score_entry_type`.
            The result is an `OrderedDict`, with entries given in descending order of score.
        """
        from chemrecon.schema.procedural_relation_entrygraph import ProceduralRelationEG

        g = entrygraph.g.copy()

        # Add source vertex and connect to initial vertices
        source_vertex_index = g.add_node(SourceVertexArtificial())
        for init_v_index in entrygraph.initial_vertices:
            g.add_edge(source_vertex_index, init_v_index, SourceEdgeArtificial())

        # Get the weight for all edges
        edge_weight_dict: dict[Edge, float] = dict()
        for e_index in g.edge_indices():
            e = g.get_edge_data_by_index(e_index)
            source_index = g.get_edge_endpoints_by_index(e_index)[0]
            if isinstance(e, Edge):
                edge_weight_dict[e] = (
                        self.relation_weight(e.relation)
                        * self.entry_weight(g.get_node_data(source_index).entry)
                )
            else:
                edge_weight_dict[e] = 1

            # Modify edges by 'score' parameter of EG-based procedural relations
            if isinstance(e, Edge):
                if isinstance(e.relation, ProceduralRelationEG):
                    edge_weight_dict[e] *= e.relation.score

        for e_index in g.edge_indices():
            e = g.get_edge_data_by_index(e_index)

        # Produce the initial score (ranking all vertices, not normalized)
        init_scoring = rustworkx.pagerank(
            g,
            alpha = self.alpha,
            weight_fn = lambda e_: edge_weight_dict.get(e_, 1),
            dangling = {source_vertex_index: 1},  # Terminal nodes should loop back to source
            personalization = {source_vertex_index: 1}  # Random traversals go back to source
        )

        # Use only eligible entries (filter by relevant entry type)
        scoring: dict[Vertex, float] = dict()
        for v_idx, score in init_scoring.items():
            v = g.get_node_data(v_idx)
            if isinstance(v, Vertex) and isinstance(v.entry, self.score_entry_type):
                scoring[v] = score

        # Apply decay based on the generation
        if self.decay_factor != 0:
            for v, score in scoring.items():
                v: Vertex
                scoring[v] = score * ((1 - self.decay_factor)**v.generation)

        # Normalise s.t. sum(scoring.values()) == 1.0
        score_sum = sum(s for _, s in scoring.items())
        for v, score in scoring.items():
            scoring[v] = score / score_sum

        # Finalise
        return OrderedDict(
            (k.entry, v) for k, v in
            sorted(scoring.items(), key = lambda pair: pair[1], reverse = True)
        )
