"""Tests for graph/network analysis functions."""

import math
import pytest

from horizon._horizon import (
    Community,
    GraphEdge,
    GraphNode,
    MarketGraph,
    build_correlation_graph,
    contagion_risk,
    detect_communities,
    minimum_spanning_tree,
)


# ===========================================================================
# Helper data
# ===========================================================================

def _corr_3x3():
    """Symmetric 3x3 correlation matrix."""
    return [
        [1.0, 0.8, 0.2],
        [0.8, 1.0, 0.3],
        [0.2, 0.3, 1.0],
    ]


def _names_3():
    return ["A", "B", "C"]


def _adj_3x3():
    """Symmetric 3x3 adjacency matrix."""
    return [
        [0.0, 0.8, 0.2],
        [0.8, 0.0, 0.3],
        [0.2, 0.3, 0.0],
    ]


def _corr_4x4():
    """Symmetric 4x4 correlation matrix with two clusters."""
    return [
        [1.0, 0.9, 0.1, 0.15],
        [0.9, 1.0, 0.05, 0.1],
        [0.1, 0.05, 1.0, 0.85],
        [0.15, 0.1, 0.85, 1.0],
    ]


# ===========================================================================
# build_correlation_graph
# ===========================================================================


class TestBuildCorrelationGraph:
    def test_basic(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        assert isinstance(result, MarketGraph)

    def test_result_fields(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        assert hasattr(result, "nodes")
        assert hasattr(result, "edges")
        assert hasattr(result, "communities")
        assert hasattr(result, "mst_edges")
        assert hasattr(result, "density")

    def test_nodes_count(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        assert len(result.nodes) == 3

    def test_node_has_centrality(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        node = result.nodes[0]
        assert isinstance(node, GraphNode)
        assert hasattr(node, "name")
        assert hasattr(node, "degree")
        assert hasattr(node, "betweenness")
        assert hasattr(node, "closeness")

    def test_edges_above_threshold(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.5)
        # Only A-B (0.8) is above 0.5 threshold
        assert len(result.edges) >= 1
        for edge in result.edges:
            assert abs(edge.weight) >= 0.5

    def test_high_threshold_no_edges(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.99)
        assert len(result.edges) == 0

    def test_low_threshold_all_edges(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.0)
        # 3 choose 2 = 3 possible edges
        assert len(result.edges) == 3

    def test_density_range(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        assert 0.0 <= result.density <= 1.0

    def test_empty_names(self):
        result = build_correlation_graph([], [], 0.1)
        assert len(result.nodes) == 0
        assert len(result.edges) == 0

    def test_mismatched_matrix_raises(self):
        with pytest.raises(ValueError):
            build_correlation_graph(
                [[1.0, 0.5], [0.5, 1.0]],
                ["A", "B", "C"],
                0.1,
            )

    def test_non_square_row_raises(self):
        with pytest.raises(ValueError):
            build_correlation_graph(
                [[1.0, 0.5], [0.5]],
                ["A", "B"],
                0.1,
            )

    def test_non_finite_raises(self):
        with pytest.raises(ValueError):
            build_correlation_graph(
                [[1.0, float("nan")], [float("nan"), 1.0]],
                ["A", "B"],
                0.1,
            )

    def test_mst_edges(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        # MST has n-1 edges for n nodes
        assert len(result.mst_edges) == 2

    def test_communities_exist(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        assert len(result.communities) >= 1

    def test_node_repr(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.1)
        r = repr(result.nodes[0])
        assert "GraphNode" in r

    def test_edge_repr(self):
        result = build_correlation_graph(_corr_3x3(), _names_3(), 0.0)
        r = repr(result.edges[0])
        assert "GraphEdge" in r

    def test_four_nodes_two_clusters(self):
        names = ["A", "B", "C", "D"]
        result = build_correlation_graph(_corr_4x4(), names, 0.5)
        assert len(result.nodes) == 4
        assert len(result.mst_edges) == 3

    def test_single_node_graph(self):
        result = build_correlation_graph([[1.0]], ["X"], 0.1)
        assert len(result.nodes) == 1
        assert len(result.edges) == 0


# ===========================================================================
# minimum_spanning_tree
# ===========================================================================


class TestMinimumSpanningTree:
    def test_basic(self):
        mst = minimum_spanning_tree(_corr_3x3(), _names_3())
        assert len(mst) == 2  # n-1 edges

    def test_edge_type(self):
        mst = minimum_spanning_tree(_corr_3x3(), _names_3())
        for edge in mst:
            assert isinstance(edge, GraphEdge)

    def test_empty(self):
        mst = minimum_spanning_tree([], [])
        assert len(mst) == 0

    def test_single_node(self):
        mst = minimum_spanning_tree([[1.0]], ["A"])
        assert len(mst) == 0

    def test_two_nodes(self):
        mst = minimum_spanning_tree(
            [[1.0, 0.5], [0.5, 1.0]], ["A", "B"]
        )
        assert len(mst) == 1

    def test_mismatched_raises(self):
        with pytest.raises(ValueError):
            minimum_spanning_tree([[1.0]], ["A", "B"])

    def test_four_nodes_three_edges(self):
        names = ["A", "B", "C", "D"]
        mst = minimum_spanning_tree(_corr_4x4(), names)
        assert len(mst) == 3


# ===========================================================================
# detect_communities
# ===========================================================================


class TestDetectCommunities:
    def test_basic(self):
        comms = detect_communities(_adj_3x3(), _names_3())
        assert len(comms) >= 1

    def test_community_type(self):
        comms = detect_communities(_adj_3x3(), _names_3())
        for c in comms:
            assert isinstance(c, Community)

    def test_community_fields(self):
        comms = detect_communities(_adj_3x3(), _names_3())
        c = comms[0]
        assert hasattr(c, "id")
        assert hasattr(c, "members")
        assert hasattr(c, "modularity")

    def test_all_nodes_assigned(self):
        comms = detect_communities(_adj_3x3(), _names_3())
        all_members = []
        for c in comms:
            all_members.extend(c.members)
        assert set(all_members) == set(_names_3())

    def test_empty(self):
        comms = detect_communities([], [])
        assert len(comms) == 0

    def test_mismatched_raises(self):
        with pytest.raises(ValueError):
            detect_communities([[0.0]], ["A", "B"])

    def test_community_repr(self):
        comms = detect_communities(_adj_3x3(), _names_3())
        r = repr(comms[0])
        assert "Community" in r


# ===========================================================================
# contagion_risk
# ===========================================================================


class TestContagionRisk:
    def test_basic(self):
        result = contagion_risk(_adj_3x3(), _names_3(), "A", 0.5)
        assert len(result) == 3

    def test_returns_tuples(self):
        result = contagion_risk(_adj_3x3(), _names_3(), "A", 0.5)
        for name, impact in result:
            assert isinstance(name, str)
            assert isinstance(impact, float)
            assert impact >= 0.0

    def test_shock_node_has_impact_at_least_1(self):
        result = contagion_risk(_adj_3x3(), _names_3(), "A", 0.5)
        shock_impact = [imp for (name, imp) in result if name == "A"]
        assert len(shock_impact) == 1
        # BFS propagation can add impact from paths that loop back,
        # so the shock node's impact can exceed 1.0
        assert shock_impact[0] >= 1.0

    def test_decay_zero_no_propagation(self):
        result = contagion_risk(_adj_3x3(), _names_3(), "A", 0.0)
        for name, impact in result:
            if name != "A":
                assert impact == pytest.approx(0.0, abs=1e-10)

    def test_unknown_shock_node_raises(self):
        with pytest.raises(ValueError):
            contagion_risk(_adj_3x3(), _names_3(), "Z", 0.5)

    def test_empty(self):
        result = contagion_risk([], [], "A", 0.5)
        assert len(result) == 0

    def test_sorted_by_impact(self):
        result = contagion_risk(_adj_3x3(), _names_3(), "A", 0.5)
        impacts = [imp for (_, imp) in result]
        for i in range(len(impacts) - 1):
            assert impacts[i] >= impacts[i + 1]

    def test_mismatched_raises(self):
        with pytest.raises(ValueError):
            contagion_risk([[0.0]], ["A", "B"], "A", 0.5)

    def test_non_finite_adj_works_with_zero_weight(self):
        adj = [[0.0, 0.5], [0.5, 0.0]]
        result = contagion_risk(adj, ["A", "B"], "A", 0.5)
        assert len(result) == 2

    def test_higher_decay_more_impact(self):
        result_low = contagion_risk(_adj_3x3(), _names_3(), "A", 0.1)
        result_high = contagion_risk(_adj_3x3(), _names_3(), "A", 0.9)
        # Higher decay should propagate more
        impacts_low = {n: i for n, i in result_low}
        impacts_high = {n: i for n, i in result_high}
        for name in ["B", "C"]:
            assert impacts_high[name] >= impacts_low[name] - 1e-10
