from typing import Any, Callable, List, Tuple

import numpy as np
from ragraph.analysis import cluster, similarity
from ragraph.edge import Edge
from ragraph.graph import Graph
from ragraph.node import Node


class SimilarityAnalysis:
    """Similarity analysis of nodes based upon mutual mapping relations.

    Arguments:
        cols: List of column nodes.
        rows: List of row nodes.
        edges: List of edges from column nodes to row nodes to be used in similarity
            analysis.
        col_sim_threshold: Column similarity threshold. Values below this threshold are
            pruned from the similarity matrix and the corresponding edges are removed.
            Defaults to 0.0 (no threshold).
        row_sim_threshold: Column similarity threshold. Values below this threshold are
            pruned from the similarity matrix and the corresponding edges are removed.
            Defaults to 0.0 (no threshold).

    Class Attributes:
        similarity_kind: Edge kind for similarity edges. Defaults to ``similarity``.

    Note:
        A mapping matrix relating M column nodes to N row nodes is used as input for
        the similarity analysis.
    """

    similarity_kind = "similarity"

    def __init__(
        self,
        rows: List[Node],
        cols: List[Node],
        edges: List[Edge],
        col_sim_threshold: float = 0.0,
        row_sim_threshold: float = 0.0,
    ):
        self._cols: List[Node] = []
        self._rows: List[Node] = []
        self._edges: List[Edge] = []
        self._col_sim_threshold = 0.0
        self._row_sim_threshold = 0.0
        self._graph: Graph = None  # type: ignore

        self.cols = cols
        self.rows = rows
        self.edges = edges
        self.col_sim_threshold = col_sim_threshold
        self.row_sim_threshold = row_sim_threshold

    @property
    def rows(self) -> List[Node]:
        """List of rows nodes."""
        return self._rows

    @rows.setter
    def rows(self, value: List[Node]):
        kinds = set([n.kind for n in value])
        if len(kinds) > 1:
            raise ValueError(
                f"""All row nodes must be of the same kind. Found node kinds:
                 {kinds}."""
            )

        if self._rows:
            self._rows = value
            self.update_graph()
        else:
            self._rows = value

    @property
    def cols(self) -> List[Node]:
        """List of column nodes."""
        return self._cols

    @cols.setter
    def cols(self, value: List[Node]):
        kinds = set([n.kind for n in value])
        if len(kinds) > 1:
            raise ValueError(
                f"""All columns nodes must be of the same kind. Found node kinds:
                 {kinds}."""
            )

        if self._cols:
            self._cols = value
            self.update_graph()
        else:
            self._cols = value

    @property
    def edges(self) -> List[Edge]:
        """List of edges."""
        return self._edges

    @edges.setter
    def edges(self, value: List[Edge]):
        if self._edges:
            self._edges = value
            self.update_graph()
        else:
            self._edges = value

    @property
    def row_sim_threshold(self) -> float:
        """Similarity threshold. Values below this threshold are pruned from the
        row similarity matrix and the corresponding edges are removed.
        """
        return self._row_sim_threshold

    @row_sim_threshold.setter
    def row_sim_threshold(self, value: float):
        if value < 0.0 or value >= 1.0:
            raise ValueError(
                f"""Similarity threshold must have a value between 0.0 and 1.0. Found
                value: {value}."""
            )

        if self._row_sim_threshold is not None:
            self._row_sim_threshold = float(value)
            self.update_row_similarity()
        else:
            self._row_sim_threshold = float(value)

    @property
    def col_sim_threshold(self) -> float:
        """Similarity threshold. Values below this threshold are pruned from the
        column similarity matrix and the corresponding edges are removed.
        """
        return self._col_sim_threshold

    @col_sim_threshold.setter
    def col_sim_threshold(self, value: float):
        if value < 0.0 or value >= 1.0:
            raise ValueError(
                f"""Similarity threshold must have a value between 0.0 and 1.0. Found
                value: {value}."""
            )

        if self._col_sim_threshold is not None:
            self._col_sim_threshold = float(value)
            self.update_col_similarity()
        else:
            self._col_sim_threshold = float(value)

    @property
    def graph(self) -> Graph:
        """Graph containing similarity edges."""
        if self._graph:
            return self._graph
        else:
            self.update_graph()
            return self._graph

    @property
    def row_similarity_matrix(self) -> np.ndarray:
        """The row similarity matrix based on their mapping row."""
        jaccard = similarity.jaccard_matrix(self.rows, self.row_mapping)
        jaccard[jaccard < self.row_sim_threshold] = 0.0
        return jaccard

    @property
    def col_similarity_matrix(self) -> np.ndarray:
        """The column similarity matrix based on their mapping column."""
        jaccard = similarity.jaccard_matrix(self.cols, self.col_mapping)
        jaccard[jaccard < self.col_sim_threshold] = 0.0
        return jaccard

    def update_graph(self) -> None:
        """Update Internal similarity graph"""
        self._graph = Graph(nodes=self.cols + self.rows, edges=self.edges)
        self.update_col_similarity()
        self.update_row_similarity()

    def update_row_similarity(self) -> None:
        """Update Jaccard Row Similarity Index edges between (clustered) rows."""
        self._update_similarity(self.rows, self.row_similarity_matrix)

    def update_col_similarity(self) -> None:
        """Update Jaccard Column Similarity Index edges between (clustered) columns."""
        self._update_similarity(self.cols, self.col_similarity_matrix)

    def _update_similarity(self, nodes: List[Node], mat: np.ndarray) -> None:
        """Update Jaccard Similarity Index edges between (clustered) nodes."""
        if not self.graph:
            self.update_graph()

        for e in [
            edge
            for edge in self.graph.edges_between_all(nodes, nodes)
            if edge.kind == self.similarity_kind
        ]:
            self.graph.del_edge(e)

        for row, target in enumerate(nodes):
            for col, source in enumerate(nodes):
                if row == col:
                    continue
                sim = mat[row][col]
                if not sim:
                    continue
                self.graph.add_edge(
                    Edge(
                        source,
                        target,
                        kind=self.similarity_kind,
                        weights=dict(similarity=mat[row][col]),
                    )
                )

    def cluster_rows(
        self,
        algo: Callable[[Graph, Any], Tuple[List[Node]]] = cluster.markov,
        **algo_args,
    ) -> None:
        """Cluster column nodes based on their similarity. Updates Graph in-place.

        Arguments:
            algo: Clustering algorithm. Should take a graph as first argument and
                cluster it in-place. Defaults to :obj:`cluster.markov`.
            **algo_args: Algorithm arguments. See :obj:`cluster.markov` for sensible
                defaults.
        """
        self._cluster(self.rows, algo, **algo_args)

    def cluster_cols(
        self,
        algo: Callable[[Graph, Any], Tuple[List[Node]]] = cluster.markov,
        **algo_args,
    ) -> None:
        """Cluster column nodes based on their similarity. Updates Graph in-place.

        Arguments:
            algo: Clustering algorithm. Should take a graph as first argument and
                cluster it in-place. Defaults to :obj:`cluster.markov`.
            **algo_args: Algorithm arguments. See :obj:`cluster.markov` for sensible
                defaults.
        """
        self._cluster(self.cols, algo, **algo_args)

    def _cluster(
        self,
        leafs: List[Node],
        algo: Callable[[Graph, Any], Tuple[List[Node]]] = cluster.markov,
        **algo_args,
    ) -> None:
        """Cluster column nodes based on their similarity. Updates Graph in-place.

        Arguments:
            leafs: List of row or column nodes to be clustered.
            algo: Clustering algorithm. Should take a graph as first argument and
                cluster it in-place. Defaults to :obj:`cluster.markov`.
            **algo_args: Algorithm arguments. See :obj:`cluster.markov` for sensible
                defaults.
        """
        algo(self.graph, leafs=leafs, **algo_args, inplace=True)  # type: ignore

    def row_mapping(self, row: Node) -> List[bool]:
        """Boolean possession checklist for a row node w.r.t. :obj:`self.cols`."""
        return [self.check_mapping(col, row) for col in self.cols]

    def col_mapping(self, col: Node) -> List[bool]:
        """Boolean possession checklist for a column node w.r.t. :obj:`self.rows`."""
        return [self.check_mapping(col, row) for row in self.rows]

    def check_mapping(self, col: Node, row: Node) -> bool:
        """Check whether a column node maps to a row node."""
        if len(self.graph.directed_edges[col.name][row.name]) > 0:
            return True
        else:
            return False
