from __future__ import annotations

from collections.abc import Generator, Hashable

import igraph  # type: ignore[import-untyped]
import networkx as nx
import pandas as pd

from graph.exploration.algorithms.graph_algorithm_registry import GRAPH_FEATURE_ALGORITHMS_BY_NAME
from graph.exploration.algorithms.graph_features_params import GraphFeaturesParams, selected_graph_feature_names
from graph.exploration.algorithms.in_memory.projected_graph_loader import LoadedProjectedGraph

FeatureRow = dict[str, object]
FeatureRowsByKey = dict[Hashable, FeatureRow]

STRING_DTYPE = "object"

def canonical_edge_key(
    source_key: Hashable,
    target_key: Hashable,
    *,
    directed: bool,
) -> Hashable:
    if directed:
        return (source_key, target_key)
    return frozenset((source_key, target_key))


def iter_simple_edges(
    graph_data: LoadedProjectedGraph,
    *,
    directed: bool,
) -> Generator[tuple[Hashable, Hashable], None, None]:
    seen_edges: set[Hashable] = set()

    for edge in graph_data.edges:
        edge_key = canonical_edge_key(edge.source_key, edge.target_key, directed=directed)
        if edge_key in seen_edges:
            continue

        seen_edges.add(edge_key)
        yield edge.source_key, edge.target_key


def iter_edges(graph_data: LoadedProjectedGraph) -> Generator[tuple[Hashable, Hashable], None, None]:
    for edge in graph_data.edges:
        yield edge.source_key, edge.target_key


def build_networkx_graph(graph_data: LoadedProjectedGraph, *, directed: bool) -> nx.Graph | nx.DiGraph:
    graph: nx.Graph | nx.DiGraph = nx.DiGraph() if directed else nx.Graph()

    for node in graph_data.nodes.values():
        graph.add_node(node.key, node_id=node.node_id, node_group=node.node_group)

    for source_key, target_key in iter_simple_edges(graph_data, directed=directed):
        graph.add_edge(source_key, target_key)

    return graph


def build_igraph_graph(
    graph_data: LoadedProjectedGraph,
    *,
    directed: bool,
) -> tuple[igraph.Graph, list[Hashable]]:
    ordered_nodes = list(graph_data.nodes)
    node_index = {node_key: index for index, node_key in enumerate(ordered_nodes)}
    graph = igraph.Graph(
        n=len(ordered_nodes),
        edges=[
            (node_index[source_key], node_index[target_key])
            for source_key, target_key in iter_edges(graph_data)
        ],
        directed=directed,
    )
    return graph, ordered_nodes


def compute_degree_by_key(graph_data: LoadedProjectedGraph, *, directed: bool) -> dict[Hashable, int]:
    degrees = {node_key: 0 for node_key in graph_data.nodes}

    for source_key, target_key in iter_edges(graph_data):
        degrees[source_key] += 1
        degrees[target_key] += 1

    return degrees


def initialize_node_rows(graph_data: LoadedProjectedGraph) -> tuple[list[FeatureRow], FeatureRowsByKey]:
    rows = [
        {
            "_node_key": node.key,
            "node_id": node.node_id,
            "node_group": node.node_group,
        }
        for node in graph_data.nodes.values()
    ]
    return rows, {row["_node_key"]: row for row in rows}


def _selected_output_columns(params: GraphFeaturesParams) -> list[tuple[str, str]]:
    columns: list[tuple[str, str]] = []
    for feature_name in selected_graph_feature_names(params):
        columns.extend(GRAPH_FEATURE_ALGORITHMS_BY_NAME[feature_name].output_columns)
    return columns


def build_empty_output_frame(params: GraphFeaturesParams) -> pd.DataFrame:
    columns: list[tuple[str, str]]
    if params.output_type == "output_edges":
        columns = [
            ("source_id", STRING_DTYPE),
            ("source_group", STRING_DTYPE),
            ("target_id", STRING_DTYPE),
            ("target_group", STRING_DTYPE),
            ("edge_group", STRING_DTYPE),
        ]
        for column_name, dtype in _selected_output_columns(params):
            columns.append((f"{column_name}_source", dtype))
            columns.append((f"{column_name}_target", dtype))
    else:
        columns = [
            ("node_id", STRING_DTYPE),
            ("node_group", STRING_DTYPE),
            *_selected_output_columns(params),
        ]

    return pd.DataFrame({column_name: pd.Series(dtype=dtype) for column_name, dtype in columns})


def iter_node_batches(
    node_rows: list[FeatureRow],
    batch_size: int,
) -> Generator[pd.DataFrame, None, None]:
    cleaned_rows = [{key: value for key, value in row.items() if key != "_node_key"} for row in node_rows]
    cleaned_rows.sort(key=lambda row: (str(row["node_group"]), str(row["node_id"])))

    for offset in range(0, len(cleaned_rows), batch_size):
        yield pd.DataFrame(cleaned_rows[offset : offset + batch_size])


def iter_edge_batches(
    graph_data: LoadedProjectedGraph,
    node_rows: list[FeatureRow],
    batch_size: int,
) -> Generator[pd.DataFrame, None, None]:
    feature_rows_by_key = {row["_node_key"]: row for row in node_rows}
    rows = []
    for edge in graph_data.edges:
        row: FeatureRow = {
            "source_id": edge.source_id,
            "source_group": edge.source_group,
            "target_id": edge.target_id,
            "target_group": edge.target_group,
            "edge_group": edge.edge_group,
        }
        for suffix, node_key in (("source", edge.source_key), ("target", edge.target_key)):
            feature_row = feature_rows_by_key[node_key]
            for key, value in feature_row.items():
                if key == "_node_key":
                    continue
                if key in {"node_id", "node_group"}:
                    continue
                row[f"{key}_{suffix}"] = value
        rows.append(row)

    for offset in range(0, len(rows), batch_size):
        yield pd.DataFrame(rows[offset : offset + batch_size])
