"""Tests for conditional independence testing."""

from __future__ import annotations

import random

from orxaq.intelligence.independence import PartialCorrelationTest, pearson_correlation


class TestPearsonCorrelation:
    def test_perfectly_correlated(self):
        x = [1.0, 2.0, 3.0, 4.0, 5.0]
        y = [2.0, 4.0, 6.0, 8.0, 10.0]
        r = pearson_correlation(x, y)
        assert abs(r - 1.0) < 0.01

    def test_negatively_correlated(self):
        x = [1.0, 2.0, 3.0, 4.0, 5.0]
        y = [10.0, 8.0, 6.0, 4.0, 2.0]
        r = pearson_correlation(x, y)
        assert abs(r - (-1.0)) < 0.01

    def test_independent_data(self):
        rng = random.Random(42)
        x = [rng.gauss(0, 1) for _ in range(500)]
        y = [rng.gauss(0, 1) for _ in range(500)]
        r = pearson_correlation(x, y)
        assert abs(r) < 0.15

    def test_short_data_returns_zero(self):
        assert pearson_correlation([1.0], [2.0]) == 0.0
        assert pearson_correlation([1.0, 2.0], [3.0, 4.0]) == 0.0


class TestPartialCorrelation:
    def _make_chain_data(self, n: int = 500, seed: int = 42) -> list[dict]:
        """Generate X → Z → Y chain data."""
        rng = random.Random(seed)
        data = []
        for _ in range(n):
            x = rng.gauss(0, 1)
            z = 2.0 * x + rng.gauss(0, 0.5)
            y = 1.5 * z + rng.gauss(0, 0.5)
            data.append({"X": x, "Z": z, "Y": y})
        return data

    def test_unconditional_correlation(self):
        """X and Y should be correlated (via Z) unconditionally."""
        data = self._make_chain_data()
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert not result.independent  # X and Y are dependent
        assert result.p_value < 0.05

    def test_conditioning_removes_spurious_correlation(self):
        """X and Y should be independent given Z in a chain."""
        data = self._make_chain_data(n=1000)
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y", frozenset({"Z"}))
        assert result.independent  # X ⊥ Y | Z
        assert result.p_value > 0.05

    def test_significant_correlation_low_pvalue(self):
        """Strongly correlated data should have very low p-value."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            x = rng.gauss(0, 1)
            y = 5.0 * x + rng.gauss(0, 0.1)
            data.append({"X": x, "Y": y})

        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert result.p_value < 0.001
        assert not result.independent

    def test_independent_high_pvalue(self):
        """Independent data should have high p-value."""
        rng = random.Random(42)
        data = []
        for _ in range(200):
            data.append({"X": rng.gauss(0, 1), "Y": rng.gauss(0, 1)})

        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert result.p_value > 0.05
        assert result.independent

    def test_method_field(self):
        data = [{"X": float(i), "Y": float(i)} for i in range(20)]
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert result.method == "partial_correlation"

    def test_small_data_returns_independent(self):
        """With very small data, should default to independent."""
        data = [{"X": 1.0, "Y": 2.0}]
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert result.independent
        assert result.p_value == 1.0

    def test_fork_structure(self):
        """In X ← Z → Y, X ⊥ Y | Z."""
        rng = random.Random(42)
        data = []
        for _ in range(1000):
            z = rng.gauss(0, 1)
            x = 2.0 * z + rng.gauss(0, 0.5)
            y = 1.5 * z + rng.gauss(0, 0.5)
            data.append({"X": x, "Z": z, "Y": y})

        ci = PartialCorrelationTest()
        # Unconditional: X and Y dependent (via Z)
        result_uncond = ci.test(data, "X", "Y")
        assert not result_uncond.independent

        # Conditional on Z: X ⊥ Y
        result_cond = ci.test(data, "X", "Y", frozenset({"Z"}))
        assert result_cond.independent


# ── Targeted coverage tests ──────────────────────────────────


class TestPartialCorrelationEdgeCases:
    """Cover lines 91, 95-96 (None/TypeError fallback), 123, 199."""

    def test_none_values_fallback(self):
        """None values in data should fall back to 0.0."""
        data = [{"X": None, "Y": 1.0}] * 10
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert isinstance(result.p_value, float)

    def test_non_numeric_values_fallback(self):
        """Non-numeric string values should fall back to 0.0."""
        data = [{"X": "bad", "Y": 1.0}, {"X": 2.0, "Y": [1]}] * 5
        ci = PartialCorrelationTest()
        result = ci.test(data, "X", "Y")
        assert isinstance(result.p_value, float)

    def test_precomputed_small_n(self):
        """test_precomputed with n<5 returns independent."""
        ci = PartialCorrelationTest()
        col_vals = {"X": [1.0, 2.0], "Y": [3.0, 4.0]}
        result = ci.test_precomputed(col_vals, 2, "X", "Y", frozenset())
        assert result.independent
        assert result.p_value == 1.0

    def test_fisher_z_low_df(self):
        """Fisher z-transform with df < 1 should return p=1.0."""
        ci = PartialCorrelationTest()
        # n=5, k=3 => df = 5-3-3 = -1 < 1 → p=1.0
        data = [
            {"X": float(i), "Y": float(i), "A": float(i), "B": float(i), "C": float(i)}
            for i in range(5)
        ]
        result = ci.test(data, "X", "Y", frozenset({"A", "B", "C"}))
        assert result.p_value == 1.0
