from __future__ import annotations

import copy
import sys
import typing
from collections.abc import Hashable, Iterable
from pathlib import Path
from typing import Any, Final, Literal, Self, cast, overload

import networkx as nx
import numpy as np
import numpy.typing as npt
from networkx import DiGraph, Graph
from pandas import DataFrame

from lang_main.constants import EDGE_WEIGHT_DECIMALS
from lang_main.io import load_pickle, save_pickle
from lang_main.loggers import logger_graphs as logger
from lang_main.types import (
    CytoscapeData,
    EdgeWeight,
    NodeTitle,
    WeightData,
)

# TODO change logging behaviour, add logging to file
LOGGING_DEFAULT: Final[bool] = False


def get_graph_metadata(
    graph: Graph | DiGraph,
    logging: bool = LOGGING_DEFAULT,
) -> dict[str, int]:
    # info about graph
    graph_info: dict[str, int] = {}
    # nodes and edges
    num_nodes = len(graph.nodes)
    num_edges = len(graph.edges)
    # edge weights
    min_edge_weight: int = 1_000_000
    max_edge_weight: int = 0
    for edge in graph.edges:
        weight = typing.cast(int, graph[edge[0]][edge[1]]['weight'])
        if weight < min_edge_weight:
            min_edge_weight = weight
        if weight > max_edge_weight:
            max_edge_weight = weight

    # memory
    edge_mem = sum([sys.getsizeof(e) for e in graph.edges])
    node_mem = sum([sys.getsizeof(n) for n in graph.nodes])
    total_mem = edge_mem + node_mem

    graph_info.update(
        num_nodes=num_nodes,
        num_edges=num_edges,
        min_edge_weight=min_edge_weight,
        max_edge_weight=max_edge_weight,
        node_memory=node_mem,
        edge_memory=edge_mem,
        total_memory=total_mem,
    )

    if logging:
        logger.info('Graph properties: %d Nodes, %d Edges', num_nodes, num_edges)
        logger.info('Node memory: %.2f KB', (node_mem / 1024))
        logger.info('Edge memory: %.2f KB', (edge_mem / 1024))
        logger.info('Total memory: %.2f KB', (total_mem / 1024))

    return graph_info


def update_graph(
    graph: Graph | DiGraph,
    *,
    batch: Iterable[tuple[Hashable, Hashable]] | None = None,
    parent: Hashable | None = None,
    child: Hashable | None = None,
    weight_connection: int | None = None,
) -> None:
    # !! not necessary to check for existence of nodes
    # !! feature already implemented in NetworkX ``add_edge``
    """
    # check if nodes already in Graph
    if parent not in graph:
        graph.add_node(parent)
    if child not in graph:
        graph.add_node(child)
    """
    if weight_connection is None:
        weight_connection = 1
    # check if edge not in Graph
    if batch is not None:
        graph.add_edges_from(batch, weight=weight_connection)
    elif not graph.has_edge(parent, child):
        # create new edge, nodes will be created if not already present
        graph.add_edge(parent, child, weight=weight_connection)
    else:
        # update edge
        weight = graph[parent][child]['weight']
        weight += weight_connection
        graph[parent][child]['weight'] = weight


# build undirected adjacency matrix
def convert_graph_to_undirected(
    graph: DiGraph,
    logging: bool = LOGGING_DEFAULT,
    cast_int: bool = False,
) -> Graph:
    dtype = np.float32
    if cast_int:
        dtype = np.uint32
    # get adjacency matrix
    adj_mat = typing.cast(DataFrame, nx.to_pandas_adjacency(G=graph, dtype=dtype))
    arr = typing.cast(npt.NDArray[np.float32 | np.uint32], adj_mat.to_numpy())
    if not cast_int:
        arr = arr * (10**EDGE_WEIGHT_DECIMALS)
        arr = np.round(arr, decimals=0)
        arr = arr.astype(np.uint32)
    # build undirected array: adding edges of lower triangular matrix to upper one
    arr_upper = np.triu(arr)
    arr_lower = np.tril(arr)
    arr_lower = np.rot90(np.fliplr(arr_lower))
    arr_new = arr_upper + arr_lower
    if not cast_int:
        arr_new = (arr_new / 10**EDGE_WEIGHT_DECIMALS).astype(np.float32)
        arr_new = np.round(arr_new, decimals=EDGE_WEIGHT_DECIMALS)
    # assign new data and create graph
    adj_mat.loc[:] = arr_new  # type: ignore
    graph_undir = typing.cast(Graph, nx.from_pandas_adjacency(df=adj_mat))

    # info about graph
    if logging:
        logger.info('Successfully converted graph to one with undirected edges.')
    _ = get_graph_metadata(graph=graph_undir, logging=logging)

    return graph_undir


def convert_graph_to_cytoscape(
    graph: Graph | DiGraph,
) -> tuple[list[CytoscapeData], WeightData]:
    cyto_data: list[CytoscapeData] = []
    # iterate over nodes
    nodes = cast(Iterable[NodeTitle], graph.nodes)
    for node in nodes:
        node_data: CytoscapeData = {
            'data': {
                'id': node,
                'label': node,
            }
        }
        cyto_data.append(node_data)
    # iterate over edges
    weights: set[int] = set()

    edges = cast(
        Iterable[
            tuple[
                NodeTitle,
                NodeTitle,
                EdgeWeight,
            ]
        ],
        graph.edges.data('weight', default=1),  # type: ignore
    )
    for source, target, weight in edges:
        weights.add(weight)
        edge_data: CytoscapeData = {
            'data': {
                'source': source,
                'target': target,
                'weight': weight,
            }
        }
        cyto_data.append(edge_data)

    # TODO: add internal behaviour (if edge added check for new min/max)
    min_weight: int = 0
    max_weight: int = 0
    if weights:
        min_weight = min(weights)
        max_weight = max(weights)
    weight_metadata: WeightData = {'min': min_weight, 'max': max_weight}

    return cyto_data, weight_metadata


def filter_graph_by_edge_weight(
    graph: TokenGraph,
    bound_lower: int | None,
    bound_upper: int | None,
) -> TokenGraph:
    """filters all edges which are within the provided bounds

    Parameters
    ----------
    bound_lower : int | None
        lower bound for edge weights, edges with weight equal to this value are retained
    bound_upper : int | None
        upper bound for edge weights, edges with weight equal to this value are retained

    Returns
    -------
    TokenGraph
        a copy of the graph with filtered edges
    """
    original_graph_edges = copy.deepcopy(graph.edges)
    filtered_graph = graph.copy()

    if not any([bound_lower, bound_upper]):
        logger.warning('No bounds provided, returning original graph.')
        return filtered_graph

    for edge in original_graph_edges:
        weight = typing.cast(int, filtered_graph[edge[0]][edge[1]]['weight'])
        if bound_lower is not None and weight < bound_lower:
            filtered_graph.remove_edge(edge[0], edge[1])
        if bound_upper is not None and weight > bound_upper:
            filtered_graph.remove_edge(edge[0], edge[1])

    if filtered_graph._undirected is not None:
        filtered_graph.to_undirected(inplace=True, logging=False)

    filtered_graph.update_metadata(logging=False)

    return filtered_graph


def filter_graph_by_node_degree(
    graph: TokenGraph,
    bound_lower: int | None,
    bound_upper: int | None,
) -> TokenGraph:
    """filters all nodes which are within the provided bounds by their degree

    Parameters
    ----------
    bound_lower : int | None
        lower bound for node degree, nodes with degree equal to this value are retained
    bound_upper : int | None
        upper bound for node degree, nodes with degree equal to this value are retained

    Returns
    -------
    TokenGraph
        a copy of the graph with filtered nodes
    """
    # filter nodes by degree
    original_graph_nodes = copy.deepcopy(graph.nodes)
    filtered_graph = graph.copy()

    if not any([bound_lower, bound_upper]):
        logger.warning('No bounds provided, returning original graph.')
        return filtered_graph

    for node in original_graph_nodes:
        degree = filtered_graph.degree[node]  # type: ignore
        if bound_lower is not None and degree < bound_lower:
            filtered_graph.remove_node(node)
        if bound_upper is not None and degree > bound_upper:
            filtered_graph.remove_node(node)

    if filtered_graph._undirected is not None:
        filtered_graph.to_undirected(inplace=True, logging=False)

    filtered_graph.update_metadata(logging=False)

    return filtered_graph


def normalise_array_linear(
    array: npt.NDArray[np.float_],
) -> npt.NDArray[np.float32]:
    """apply standard linear normalisation

    Parameters
    ----------
    array : npt.NDArray[np.float_]
        array which shall be normalised

    Returns
    -------
    npt.NDArray[np.float32]
        min/max normalised array
    """
    arr_norm = (array - array.min()) / (array.max() - array.min())

    return arr_norm.astype(np.float32)


def weight_scaling(
    weights: npt.NDArray[np.float32],
    a: float = 1.1,
    b: float = 0.05,
) -> npt.NDArray[np.float32]:
    """non-linear scaling of already normalised edge weights [0;1]: bigger weights
    have smaller weight delta than smaller weights. Bigger values für parameter
    `b` reinforce this effect.
    Based on:
    https://math.stackexchange.com/questions/4297805/exponential-function-that-passes-through-0-0-and-1-1-with-variable-slope

    With default values the range of edge weights lies approximately between [0.1; 1]

    Parameters
    ----------
    weights : npt.NDArray[np.float32]
        pre-normalised edge weights as 1D array
    a : float, optional
        factor to determine the value for edge weights with value 0
        with default approx. 0.1, by default 1.1
    b : float, optional
        adjust the curvature, smaller values increase it, by default 0.05

    Returns
    -------
    npt.NDArray[np.float32]
        non-linear adjusted edge weights as 1D array
    """
    adjusted_weights = (b**weights - a) / (b - a)

    return np.round(adjusted_weights, decimals=EDGE_WEIGHT_DECIMALS)


def rescale_edge_weights(
    graph: TokenGraph,
) -> TokenGraph:
    graph = graph.copy()

    weights = cast(list[int], [data['weight'] for data in graph.edges.values()])
    w_log = cast(npt.NDArray[np.float32], np.log(weights, dtype=np.float32))
    weights_norm = normalise_array_linear(w_log)
    weights_adjusted = weight_scaling(weights_norm)
    # assign new weight values
    for idx, (node_1, node_2) in enumerate(list(graph.edges)):
        graph[node_1][node_2]['weight'] = weights_adjusted[idx]

    graph.rescaled_weights = True
    graph.update_metadata(logging=False)

    return graph


# ** ---------------------------------------
class TokenGraph(DiGraph):
    def __init__(
        self,
        name: str = 'TokenGraph',
        enable_logging: bool = True,
        incoming_graph_data: Any | None = None,
        **attr,
    ) -> None:
        super().__init__(incoming_graph_data, **attr)
        # logging of different actions
        self.logging = enable_logging
        # properties
        self._name = name
        # directed and undirected graph data
        self._directed = self
        self._metadata_directed: dict[str, int] = {}
        self._undirected: Graph | None = None
        self._metadata_undirected: dict[str, int] = {}
        # indicate rescaled weights
        self.rescaled_weights: bool = False

    def __repr__(self) -> str:
        return self.__str__()

    def __str__(self) -> str:
        return (
            f'TokenGraph(name: {self.name}, number of nodes: '
            f'{len(self.nodes)}, number of edges: '
            f'{len(self.edges)})'
        )

    def disable_logging(self) -> None:
        self.logging = False

    # !! only used to verify that saving was done correctly
    """
    def __key(self) -> tuple[Hashable, ...]:
        return (self.name, tuple(self.nodes), tuple(self.edges))
    
    def __hash__(self) -> int:
        return hash(self.__key())
    """

    def copy(self) -> Self:
        """returns a (deep) copy of the graph

        Returns
        -------
        Self
            deep copy of the graph
        """
        return copy.deepcopy(self)

    @property
    def name(self) -> str:
        return self._name

    @property
    def directed(self) -> Self:
        return self._directed

    @property
    def undirected(self) -> Graph | None:
        return self._undirected

    @property
    def metadata_directed(self) -> dict[str, int]:
        return self._metadata_directed

    @property
    def metadata_undirected(self) -> dict[str, int]:
        return self._metadata_undirected

    @overload
    def to_undirected(
        self,
        inplace: Literal[True] = ...,
        logging: bool | None = ...,
    ) -> None: ...

    @overload
    def to_undirected(
        self,
        inplace: Literal[False],
        logging: bool | None = ...,
    ) -> Graph: ...

    def to_undirected(
        self,
        inplace: bool = True,
        logging: bool | None = None,
    ) -> Graph | None:
        if logging is None:
            logging = self.logging
        # cast to integer edge weights only if edges were not rescaled previously
        cast_int: bool = True
        if self.rescaled_weights:
            cast_int = False

        self._undirected = convert_graph_to_undirected(
            graph=self,
            logging=logging,
            cast_int=cast_int,
        )
        self._metadata_undirected = get_graph_metadata(graph=self._undirected, logging=False)
        if not inplace:
            return self._undirected

    def update_metadata(
        self,
        logging: bool | None = None,
    ) -> None:
        if logging is None:
            logging = self.logging

        self._metadata_directed = get_graph_metadata(graph=self, logging=logging)
        if self._undirected is not None:
            self._metadata_undirected = get_graph_metadata(
                graph=self._undirected, logging=logging
            )

    def _save_prepare(
        self,
        path: Path,
        filename: str | None = None,
    ) -> Path:
        if filename is not None:
            saving_path = path.joinpath(f'{filename}')
        else:
            saving_path = path.joinpath(f'{self.name}')

        return saving_path

    def to_GraphML(
        self,
        path: Path,
        filename: str | None = None,
        directed: bool = False,
    ) -> None:
        """save one of the stored graphs to GraphML format on disk,

        Parameters
        ----------
        path : Path
            target path for saving the file
        filename : str | None, optional
            filename to be given, by default None
        directed : bool, optional
            indicator whether directed or undirected graph
            should be exported, by default False (undirected)

        Raises
        ------
        ValueError
            undirected graph should be exported but is not available
        """
        saving_path = self._save_prepare(path=path, filename=filename)

        if directed:
            target_graph = self._directed
        elif not directed and self._undirected is not None:
            target_graph = self._undirected
        else:
            raise ValueError('No undirected graph available.')

        saving_path = saving_path.with_suffix('.graphml')
        nx.write_graphml(G=target_graph, path=saving_path)
        logger.info('Successfully saved graph as GraphML file under %s.', saving_path)

    def to_pickle(
        self,
        path: Path,
        filename: str | None = None,
    ) -> None:
        """save whole TokenGraph object as pickle file

        Parameters
        ----------
        path : Path
            target path for saving the file
        filename : str | None, optional
            filename to be given, by default None
        """
        saving_path = self._save_prepare(path=path, filename=filename)
        saving_path = saving_path.with_suffix('.pickle')
        save_pickle(obj=self, path=saving_path)

    @classmethod
    def from_file(
        cls,
        path: Path,
    ) -> Self:
        # !! no validity checks for pickle files
        # !! GraphML files not correct because not all properties
        # !! are parsed correctly
        # TODO REWORK
        match path.suffix:
            case '.graphml':
                graph = typing.cast(Self, nx.read_graphml(path, node_type=int))
                logger.info('Successfully loaded graph from GraphML file %s.', path)
            case '.pkl' | '.pickle':
                graph = typing.cast(Self, load_pickle(path))
                logger.info('Successfully loaded graph from pickle file %s.', path)
            case _:
                raise ValueError('File format not supported.')

        return graph

    @classmethod
    def from_pickle(
        cls,
        path: str | Path,
    ) -> Self:
        if isinstance(path, str):
            path = Path(path)

        if path.suffix not in ('.pkl', '.pickle'):
            raise ValueError('File format not supported.')

        graph = typing.cast(Self, load_pickle(path))

        return graph
