"""Graph merge rules used during hierarchical summarisation."""

from __future__ import annotations

from collections import defaultdict
from typing import DefaultDict, Iterable, MutableMapping

import networkx as nx


def _normalise_attribute(value: object) -> object:
    """Ensure attribute values are hashable for consensus aggregation."""

    if isinstance(value, dict):
        return tuple(sorted((k, _normalise_attribute(v))
                     for k, v in value.items()))
    if isinstance(value, (list, tuple, set)):
        return tuple(_normalise_attribute(v) for v in value)
    return value


def hierarchical_merge(
    graphs: Iterable[nx.DiGraph],
    *,
    weights: Iterable[float] | None = None,
) -> nx.DiGraph:
    """Merge ``graphs`` while enforcing acyclicity and attribute consensus."""

    graphs = list(graphs)
    if not graphs:
        return nx.DiGraph()

    if weights is None:
        weights = [1.0] * len(graphs)

    weight_iter = list(weights)
    if len(weight_iter) != len(graphs):
        raise ValueError("weights length must match number of graphs")

    result = nx.DiGraph()

    node_attrs: DefaultDict[tuple[str, str], DefaultDict[object, float]] = defaultdict(
        lambda: defaultdict(float))
    edge_attrs: DefaultDict[
        tuple[tuple[str, str], str], DefaultDict[object, float]
    ] = defaultdict(lambda: defaultdict(float))
    edge_support: dict[tuple[str, str], float] = defaultdict(float)

    for g, weight in zip(graphs, weight_iter):
        for node, attrs in g.nodes(data=True):
            result.add_node(node)
            for key, value in attrs.items():
                normalised = _normalise_attribute(value)
                node_attrs[(node, key)][normalised] += float(weight)
        for u, v, attrs in g.edges(data=True):
            result.add_node(u)
            result.add_node(v)
            result.add_edge(u, v)
            edge_support[(u, v)] += float(weight)
            for key, value in attrs.items():
                normalised = _normalise_attribute(value)
                edge_attrs[((u, v), key)][normalised] += float(weight)

    for (node, key), counter in node_attrs.items():
        value = _consensus(counter)
        if value is not None:
            result.nodes[node][key] = value

    for ((u, v), key), counter in edge_attrs.items():
        value = _consensus(counter)
        if value is not None:
            result.edges[u, v][key] = value

    # Remove edges that participate in cycles to respect DAG requirement.
    while True:
        cycles = list(nx.simple_cycles(result))
        if not cycles:
            break

        edges_to_remove: set[tuple[str, str]] = set()
        for cycle in cycles:
            ordered_cycle = list(cycle)
            cycle_edges = list(
                zip(ordered_cycle, ordered_cycle[1:] + [ordered_cycle[0]]))
            edge_to_remove = min(
                cycle_edges,
                key=lambda edge: (edge_support.get(edge, 0.0), edge),
            )
            edges_to_remove.add(edge_to_remove)

        for edge in edges_to_remove:
            if result.has_edge(*edge):
                result.remove_edge(*edge)

    return result


def _consensus(counter: MutableMapping[object, float]) -> object | None:
    if not counter:
        return None
    value, strength = max(counter.items(), key=lambda item: item[1])
    total = float(sum(counter.values()))
    if total == 0:
        return value
    if strength / total < 0.4:
        return None
    return value
