"""Performance regression tests for causal discovery algorithms.

Validates that optimizations (memoization, pre-extraction, adjacency maps)
produce correct results and exercise the optimized code paths.
"""

from __future__ import annotations

import random

from orxaq.intelligence.consensus import ConsensusEngine
from orxaq.intelligence.ges_algorithm import GESAlgorithm
from orxaq.intelligence.independence import PartialCorrelationTest
from orxaq.intelligence.pc_algorithm import PCAlgorithm


def generate_synthetic_data(
    n_samples: int,
    n_vars: int,
    seed: int = 42,
) -> tuple[list[dict], list[str]]:
    """Generate synthetic tabular data with string variable names."""
    rng = random.Random(seed)
    variables = [f"v{i}" for i in range(n_vars)]
    data = []
    for _ in range(n_samples):
        row = {v: rng.gauss(0, 1) for v in variables}
        data.append(row)
    return data, variables


class TestPCPerformance:
    """PC algorithm performance and correctness under optimization."""

    def test_pc_200_samples_8_vars(self) -> None:
        data, variables = generate_synthetic_data(200, 8)
        pc = PCAlgorithm(alpha=0.05, max_cond_size=2)
        result = pc.fit(data, variables)
        assert result.ci_tests_run > 0
        assert isinstance(result.skeleton, frozenset)

    def test_pc_uses_precomputed_path(self) -> None:
        """Verify PC skeleton phase uses test_precomputed internally."""
        data, variables = generate_synthetic_data(100, 5)
        pc = PCAlgorithm(alpha=0.05, max_cond_size=2)
        result = pc.fit(data, variables)
        # If precomputed path works, results are valid
        assert result.ci_tests_run > 0
        for edge in result.skeleton:
            assert edge[0] in variables
            assert edge[1] in variables

    def test_pc_results_deterministic(self) -> None:
        """Same data should produce identical results."""
        data, variables = generate_synthetic_data(150, 6)
        pc = PCAlgorithm(alpha=0.05, max_cond_size=2)
        r1 = pc.fit(data, variables)
        r2 = pc.fit(data, variables)
        assert r1.skeleton == r2.skeleton
        assert r1.oriented == r2.oriented
        assert r1.ci_tests_run == r2.ci_tests_run


class TestGESPerformance:
    """GES algorithm performance with adjacency-based cycle check."""

    def test_ges_200_samples_6_vars(self) -> None:
        data, variables = generate_synthetic_data(200, 6)
        ges = GESAlgorithm()
        result = ges.fit(data, variables)
        assert isinstance(result.edges, frozenset)
        assert result.additions >= 0
        assert result.removals >= 0

    def test_ges_with_required_edges(self) -> None:
        data, variables = generate_synthetic_data(100, 5)
        required = {(variables[0], variables[1])}
        ges = GESAlgorithm()
        result = ges.fit(data, variables, required_edges=required)
        # Required edge must be in result
        assert (variables[0], variables[1]) in result.edges

    def test_ges_results_deterministic(self) -> None:
        data, variables = generate_synthetic_data(150, 6)
        ges = GESAlgorithm()
        r1 = ges.fit(data, variables)
        r2 = ges.fit(data, variables)
        assert r1.edges == r2.edges
        assert r1.bic_score == r2.bic_score


class TestConsensusPerformance:
    """Consensus engine with both PC and GES."""

    def test_consensus_200_samples(self) -> None:
        data, variables = generate_synthetic_data(200, 6)
        engine = ConsensusEngine(alpha=0.05, threshold=0.5)
        result = engine.discover(data, variables)
        assert result.algorithms_used == ("pc", "ges")
        for edge in result.edges:
            assert edge.source in variables
            assert edge.target in variables


class TestPrecomputedCI:
    """Verify test_precomputed gives identical results to test."""

    def test_precomputed_matches_standard(self) -> None:
        data, variables = generate_synthetic_data(100, 5)
        ci = PartialCorrelationTest()

        x, y = variables[0], variables[1]
        z_set = frozenset([variables[2]])

        # Standard path
        r1 = ci.test(data, x, y, z_set, alpha=0.05)

        # Pre-extracted path
        col_vals: dict[str, list[float]] = {v: [] for v in variables}
        for row in data:
            for v in variables:
                col_vals[v].append(float(row.get(v, 0.0)))
        r2 = ci.test_precomputed(col_vals, len(data), x, y, z_set, alpha=0.05)

        assert r1.statistic == r2.statistic
        assert r1.p_value == r2.p_value
        assert r1.independent == r2.independent

    def test_memoization_cache_reuse(self) -> None:
        """Verify memoization produces same results and populates cache."""
        data, variables = generate_synthetic_data(100, 5)
        ci = PartialCorrelationTest()

        col_vals: dict[str, list[float]] = {v: [] for v in variables}
        for row in data:
            for v in variables:
                col_vals[v].append(float(row.get(v, 0.0)))

        cache: dict[tuple[str, str, frozenset[str]], float] = {}
        x, y = variables[0], variables[1]
        z_set = frozenset([variables[2], variables[3]])

        r1 = ci.test_precomputed(col_vals, len(data), x, y, z_set, _cache=cache)
        # Cache should have entries after first call
        assert len(cache) > 0

        # Second call should use cache and give same result
        r2 = ci.test_precomputed(col_vals, len(data), x, y, z_set, _cache=cache)
        assert r1.statistic == r2.statistic
        assert r1.p_value == r2.p_value
