"""
Tests for measurement_engine.py — Raw Measurement Engine.

Validates:
- Embedding strategies (MIN_MAX, LINEAR_SCALE, MAX_NORM, ZSCORE_SIGMOID)
- Kernel invariant identities (F = 1 − ω, IC = exp(κ), IC ≤ F)
- τ_R computation and INF_REC sentinel handling
- Regime classification
- psi.csv and invariants.json output generation
- Casepack scaffolding
- Edge cases (constant data, single timestep, NaN handling)
- safe_tau_R / tau_R_display utilities

Cross-references:
- KERNEL_SPECIFICATION.md  Lemmas 1-5 (range, identity, AM-GM)
- frozen_contract.py       compute_kernel, compute_tau_R
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import pytest

from umcp.measurement_engine import (
    EmbeddingConfig,
    EmbeddingSpec,
    EmbeddingStrategy,
    InvariantRow,
    MeasurementEngine,
    TraceRow,
    safe_tau_R,
    tau_R_display,
)


# ============================================================================
# Fixtures
# ============================================================================
@pytest.fixture
def engine() -> MeasurementEngine:
    """Default measurement engine."""
    return MeasurementEngine(eta=0.10, H_rec=50)


@pytest.fixture
def sample_csv(tmp_path: Path) -> Path:
    """Create a simple raw measurements CSV."""
    csv_path = tmp_path / "raw.csv"
    csv_path.write_text("x_1,x_2,x_3\n9.9,9.9,9.9\n9.8,9.7,9.6\n9.5,9.4,9.3\n9.2,9.1,9.0\n8.9,8.8,8.7\n")
    return csv_path


@pytest.fixture
def sample_array() -> np.ndarray:
    """Simple 5×3 raw measurement array."""
    return np.array(
        [
            [9.9, 9.9, 9.9],
            [9.8, 9.7, 9.6],
            [9.5, 9.4, 9.3],
            [9.2, 9.1, 9.0],
            [8.9, 8.8, 8.7],
        ]
    )


# ============================================================================
# Embedding Strategy Tests
# ============================================================================
class TestEmbedding:
    """Test embedding/normalization strategies."""

    def test_min_max_range(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """MIN_MAX should produce values in [0, 1]."""
        result = engine.from_array(sample_array)
        coords = result.coordinates_array
        assert np.all(coords >= 0.0)
        assert np.all(coords <= 1.0)

    def test_min_max_endpoints(self, engine: MeasurementEngine) -> None:
        """MIN_MAX maps min→0, max→1."""
        data = np.array([[0.0], [5.0], [10.0]])
        result = engine.from_array(data)
        coords = result.coordinates_array
        # After ε-clipping, values won't be exactly 0 or 1
        assert coords[0, 0] < 0.01  # near 0
        assert coords[2, 0] > 0.99  # near 1

    def test_linear_scale(self, engine: MeasurementEngine) -> None:
        """LINEAR_SCALE with explicit bounds."""
        data = np.array([[5.0, 7.5], [10.0, 10.0]])
        specs = [
            EmbeddingSpec(strategy=EmbeddingStrategy.LINEAR_SCALE, input_range=(0.0, 10.0)),
            EmbeddingSpec(strategy=EmbeddingStrategy.LINEAR_SCALE, input_range=(0.0, 10.0)),
        ]
        config = EmbeddingConfig(specs=specs)
        result = engine.from_array(data, embedding=config)
        coords = result.coordinates_array
        # 5/10 = 0.5, 7.5/10 = 0.75
        assert abs(coords[0, 0] - 0.5) < 0.01
        assert abs(coords[0, 1] - 0.75) < 0.01

    def test_max_norm(self, engine: MeasurementEngine) -> None:
        """MAX_NORM divides by max absolute value."""
        data = np.array([[5.0], [10.0], [8.0]])
        config = EmbeddingConfig(default_strategy=EmbeddingStrategy.MAX_NORM)
        result = engine.from_array(data, embedding=config)
        coords = result.coordinates_array
        assert abs(coords[1, 0] - 1.0) < 0.01  # max → 1 (before ε clip)
        assert abs(coords[0, 0] - 0.5) < 0.01  # 5/10 = 0.5

    def test_zscore_sigmoid(self, engine: MeasurementEngine) -> None:
        """ZSCORE_SIGMOID produces values in (0, 1)."""
        data = np.array([[1.0], [2.0], [3.0], [4.0], [5.0]])
        config = EmbeddingConfig(default_strategy=EmbeddingStrategy.ZSCORE_SIGMOID)
        result = engine.from_array(data, embedding=config)
        coords = result.coordinates_array
        assert np.all(coords > 0.0)
        assert np.all(coords < 1.0)
        # Sigmoid of z=0 (mean) should be ~0.5
        assert abs(coords[2, 0] - 0.5) < 0.01

    def test_constant_column(self, engine: MeasurementEngine) -> None:
        """Constant column should embed to 0.5."""
        data = np.array([[7.0, 5.0], [7.0, 10.0], [7.0, 15.0]])
        result = engine.from_array(data)
        coords = result.coordinates_array
        assert abs(coords[0, 0] - 0.5) < 0.01  # constant → 0.5


# ============================================================================
# Kernel Identity Tests (Tier-1)
# ============================================================================
class TestKernelIdentities:
    """Test that Tier-1 identities hold in engine output."""

    def test_F_equals_1_minus_omega(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """Lemma: F = 1 − ω."""
        result = engine.from_array(sample_array)
        for inv in result.invariants:
            assert abs(inv.F - (1 - inv.omega)) < 1e-12

    def test_IC_equals_exp_kappa(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """Lemma 3: IC = exp(κ)."""
        result = engine.from_array(sample_array)
        for inv in result.invariants:
            assert abs(inv.IC - math.exp(inv.kappa)) < 1e-12

    def test_AM_GM_inequality(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """Lemma 4: F ≥ IC (AM-GM)."""
        result = engine.from_array(sample_array)
        for inv in result.invariants:
            assert inv.F >= inv.IC - 1e-15  # tolerance for floating point

    def test_F_in_unit_interval(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """Lemma 1: F ∈ [0, 1]."""
        result = engine.from_array(sample_array)
        for inv in result.invariants:
            assert 0 <= inv.F <= 1

    def test_omega_in_unit_interval(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        """Lemma 1: ω ∈ [0, 1]."""
        result = engine.from_array(sample_array)
        for inv in result.invariants:
            assert 0 <= inv.omega <= 1


# ============================================================================
# τ_R and Return Detection Tests
# ============================================================================
class TestTauR:
    """Test τ_R computation and INF_REC handling."""

    def test_first_timestep_inf(self, engine: MeasurementEngine) -> None:
        """First timestep should have τ_R = ∞ (no history)."""
        data = np.array([[0.5, 0.5], [0.6, 0.6]])
        result = engine.from_array(data)
        assert math.isinf(result.invariants[0].tau_R)

    def test_identical_points_return(self) -> None:
        """Identical consecutive points should have finite τ_R."""
        engine = MeasurementEngine(eta=0.10, H_rec=50)
        data = np.array([[0.5, 0.5], [0.5, 0.5], [0.5, 0.5]])
        result = engine.from_array(data)
        # Second and third timesteps should find return (distance = 0 < η)
        assert math.isfinite(result.invariants[1].tau_R)
        assert math.isfinite(result.invariants[2].tau_R)

    def test_diverging_no_return(self) -> None:
        """Strongly diverging trajectory should produce INF_REC."""
        engine = MeasurementEngine(eta=0.01, H_rec=2)
        data = np.array([[0.1, 0.1], [0.5, 0.5], [0.9, 0.9]])
        result = engine.from_array(data)
        # Distance 0.1→0.5 = ~0.57 >> η=0.01 with H_rec=2
        assert math.isinf(result.invariants[2].tau_R)


# ============================================================================
# INF_REC Utility Tests
# ============================================================================
class TestINFRECUtilities:
    """Test safe_tau_R and tau_R_display."""

    @pytest.mark.parametrize(
        "value,expected",
        [
            ("INF_REC", float("inf")),
            ("inf", float("inf")),
            ("∞", float("inf")),
            ("INFINITY", float("inf")),
            (None, float("inf")),
            (float("nan"), float("inf")),
            (float("inf"), float("inf")),
            (5.0, 5.0),
            (5, 5.0),
            ("5.0", 5.0),
            ("3.14", 3.14),
            ("", float("inf")),
            ("NONE", float("inf")),
        ],
    )
    def test_safe_tau_R(self, value: object, expected: float) -> None:
        result = safe_tau_R(value)
        if math.isinf(expected):
            assert math.isinf(result)
        else:
            assert abs(result - expected) < 1e-10

    def test_tau_R_display_inf(self) -> None:
        assert tau_R_display(float("inf")) == "INF_REC"

    def test_tau_R_display_finite(self) -> None:
        assert tau_R_display(5.0) == "5"

    def test_tau_R_display_string_inf(self) -> None:
        assert tau_R_display("INF_REC") == "INF_REC"

    def test_tau_R_display_none(self) -> None:
        assert tau_R_display(None) == "INF_REC"


# ============================================================================
# CSV Input Tests
# ============================================================================
class TestCSVInput:
    """Test from_csv entry point."""

    def test_load_csv(self, engine: MeasurementEngine, sample_csv: Path) -> None:
        result = engine.from_csv(sample_csv)
        assert result.n_timesteps == 5
        assert result.n_dims == 3
        assert len(result.trace) == 5
        assert len(result.invariants) == 5

    def test_csv_with_weights(self, engine: MeasurementEngine, sample_csv: Path) -> None:
        result = engine.from_csv(sample_csv, weights=[0.5, 0.3, 0.2])
        assert result.weights == pytest.approx([0.5, 0.3, 0.2])

    def test_csv_uniform_weights(self, engine: MeasurementEngine, sample_csv: Path) -> None:
        result = engine.from_csv(sample_csv)
        assert result.weights == pytest.approx([1 / 3, 1 / 3, 1 / 3], abs=1e-6)

    def test_csv_with_t_column(self, tmp_path: Path, engine: MeasurementEngine) -> None:
        """CSV with explicit 't' column."""
        csv_path = tmp_path / "with_t.csv"
        csv_path.write_text("t,x\n0,5.0\n10,7.0\n20,9.0\n")
        result = engine.from_csv(csv_path)
        assert result.trace[0].t == 0
        assert result.trace[1].t == 10
        assert result.trace[2].t == 20


# ============================================================================
# Array Input Tests
# ============================================================================
class TestArrayInput:
    """Test from_array entry point."""

    def test_1d_array(self, engine: MeasurementEngine) -> None:
        """1-D array should be treated as single dimension."""
        data = np.array([1.0, 2.0, 3.0])
        result = engine.from_array(data)
        assert result.n_dims == 1
        assert result.n_timesteps == 3

    def test_weight_mismatch_raises(self, engine: MeasurementEngine) -> None:
        data = np.array([[1.0, 2.0], [3.0, 4.0]])
        with pytest.raises(ValueError, match="Weight dimension mismatch"):
            engine.from_array(data, weights=[0.5, 0.3, 0.2])

    def test_custom_time_indices(self, engine: MeasurementEngine) -> None:
        data = np.array([[0.5, 0.5], [0.6, 0.6]])
        result = engine.from_array(data, time_indices=[100, 200])
        assert result.trace[0].t == 100
        assert result.trace[1].t == 200


# ============================================================================
# Output Writer Tests
# ============================================================================
class TestWriters:
    """Test psi.csv and invariants.json writers."""

    def test_write_psi_csv(self, engine: MeasurementEngine, sample_array: np.ndarray, tmp_path: Path) -> None:
        result = engine.from_array(sample_array)
        psi_path = engine.write_psi_csv(result, tmp_path / "psi.csv")
        assert psi_path.exists()
        lines = psi_path.read_text().strip().split("\n")
        assert len(lines) == 6  # header + 5 rows
        assert "c_1" in lines[0]
        assert "oor_1" in lines[0]

    def test_write_invariants_json(self, engine: MeasurementEngine, sample_array: np.ndarray, tmp_path: Path) -> None:
        result = engine.from_array(sample_array)
        inv_path = engine.write_invariants_json(result, tmp_path / "invariants.json")
        assert inv_path.exists()
        doc = json.loads(inv_path.read_text())
        assert doc["schema"] == "schemas/invariants.schema.json"
        assert doc["format"] == "tier1_invariants"
        assert len(doc["rows"]) == 5
        # Check first row
        row = doc["rows"][0]
        assert "omega" in row
        assert "F" in row
        assert "IC" in row
        assert "regime" in row
        assert "label" in row["regime"]

    def test_invariants_json_inf_rec_serialization(self, engine: MeasurementEngine, tmp_path: Path) -> None:
        """τ_R = ∞ should serialize as 'INF_REC' string in JSON."""
        data = np.array([[0.5, 0.5]])
        result = engine.from_array(data)
        inv_path = engine.write_invariants_json(result, tmp_path / "inv.json")
        doc = json.loads(inv_path.read_text())
        # First (only) timestep has no history → τ_R = INF_REC
        assert doc["rows"][0]["tau_R"] == "INF_REC"


# ============================================================================
# Casepack Generation Tests
# ============================================================================
class TestCasepackGeneration:
    """Test full casepack scaffolding."""

    def test_generate_casepack(self, engine: MeasurementEngine, sample_array: np.ndarray, tmp_path: Path) -> None:
        result = engine.from_array(sample_array)
        cp_dir = tmp_path / "test_casepack"
        engine.generate_casepack(result, cp_dir, casepack_id="test_cp")

        assert (cp_dir / "manifest.json").exists()
        assert (cp_dir / "expected" / "psi.csv").exists()
        assert (cp_dir / "expected" / "invariants.json").exists()

        manifest = json.loads((cp_dir / "manifest.json").read_text())
        assert manifest["casepack"]["id"] == "test_cp"
        assert manifest["refs"]["contract"]["id"] == "UMA.INTSTACK.v1"

    def test_casepack_default_id(self, engine: MeasurementEngine, sample_array: np.ndarray, tmp_path: Path) -> None:
        result = engine.from_array(sample_array)
        cp_dir = tmp_path / "my_pack"
        engine.generate_casepack(result, cp_dir)
        manifest = json.loads((cp_dir / "manifest.json").read_text())
        assert manifest["casepack"]["id"] == "my_pack"


# ============================================================================
# Regime Classification Tests
# ============================================================================
class TestRegimeClassification:
    """Test regime classification through the engine."""

    def test_stable_regime(self, engine: MeasurementEngine) -> None:
        """High-quality coordinates should produce STABLE."""
        # Use LINEAR_SCALE with identity bounds so values pass through unchanged
        data = np.array([[0.95, 0.95, 0.95]] * 5)
        config = EmbeddingConfig(
            specs=[EmbeddingSpec(strategy=EmbeddingStrategy.LINEAR_SCALE, input_range=(0.0, 1.0))] * 3
        )
        result = engine.from_array(data, embedding=config)
        # All coordinates near 1.0 → high F, low ω → STABLE
        # ω ≈ 0.05 < 0.038 threshold may not hold due to ε-clipping, so
        # accept STABLE or WATCH (both valid for marginal ω)
        assert result.final_regime in ("STABLE", "WATCH")

    def test_degrading_regime(self) -> None:
        """Degrading data should transition toward WATCH or COLLAPSE."""
        engine = MeasurementEngine(eta=0.10, H_rec=50)
        data = np.array(
            [
                [0.95, 0.95, 0.95],
                [0.80, 0.80, 0.80],
                [0.60, 0.60, 0.60],
                [0.40, 0.40, 0.40],
                [0.20, 0.20, 0.20],
            ]
        )
        config = EmbeddingConfig(
            specs=[EmbeddingSpec(strategy=EmbeddingStrategy.LINEAR_SCALE, input_range=(0.0, 1.0))] * 3
        )
        result = engine.from_array(data, embedding=config)
        # Later timesteps should be in WATCH or COLLAPSE
        assert result.invariants[-1].regime in ("WATCH", "COLLAPSE", "CRITICAL")


# ============================================================================
# EngineResult Tests
# ============================================================================
class TestEngineResult:
    """Test EngineResult container."""

    def test_summary(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        result = engine.from_array(sample_array)
        summary = result.summary()
        assert summary["n_timesteps"] == 5
        assert summary["n_dims"] == 3
        assert "final_regime" in summary
        assert "regime_counts" in summary
        assert "omega_range" in summary
        assert "IC_range" in summary

    def test_coordinates_array(self, engine: MeasurementEngine, sample_array: np.ndarray) -> None:
        result = engine.from_array(sample_array)
        coords = result.coordinates_array
        assert coords.shape == (5, 3)
        assert np.all(coords >= 0.0)
        assert np.all(coords <= 1.0)


# ============================================================================
# Edge Cases
# ============================================================================
class TestEdgeCases:
    """Test edge cases and robustness."""

    def test_single_timestep(self, engine: MeasurementEngine) -> None:
        """Single row should work (τ_R = ∞)."""
        data = np.array([[5.0, 7.0, 9.0]])
        result = engine.from_array(data)
        assert result.n_timesteps == 1
        assert math.isinf(result.invariants[0].tau_R)

    def test_nan_handling(self, engine: MeasurementEngine) -> None:
        """NaN values should be imputed and flagged."""
        data = np.array([[5.0, np.nan, 9.0], [6.0, 7.0, 8.0]])
        result = engine.from_array(data)
        assert result.n_timesteps == 2
        # NaN should be flagged as missing
        assert result.trace[0].miss[1] is True
        assert result.trace[0].miss[0] is False

    def test_single_dimension(self, engine: MeasurementEngine) -> None:
        """Single dimension should work."""
        data = np.array([[1.0], [2.0], [3.0]])
        result = engine.from_array(data)
        assert result.n_dims == 1
        assert len(result.invariants) == 3

    def test_empty_csv_raises(self, engine: MeasurementEngine, tmp_path: Path) -> None:
        """Empty CSV should raise ValueError."""
        csv_path = tmp_path / "empty.csv"
        csv_path.write_text("x_1,x_2\n")
        with pytest.raises(ValueError, match="empty"):
            engine.from_csv(csv_path)

    def test_trace_row_to_dict(self) -> None:
        """TraceRow serialization."""
        row = TraceRow(t=0, c=[0.5, 0.6], oor=[False, True], miss=[False, False])
        d = row.to_csv_wide_dict()
        assert d["t"] == 0
        assert d["c_1"] == 0.5
        assert d["c_2"] == 0.6
        assert d["oor_1"] is False
        assert d["oor_2"] is True

    def test_invariant_row_to_dict(self) -> None:
        """InvariantRow serialization with INF_REC."""
        row = InvariantRow(
            t=0,
            omega=0.01,
            F=0.99,
            S=0.05,
            C=0.0,
            tau_R=float("inf"),
            kappa=-0.01,
            IC=0.99,
            regime="STABLE",
            critical_overlay=False,
        )
        d = row.to_dict()
        assert d["tau_R"] == "INF_REC"
        assert d["regime"]["label"] == "Stable"
        assert d["regime"]["critical_overlay"] is False


# ============================================================================
# End-to-end CSV processing (synthetic 5-timestep, 3-dim data)
# ============================================================================
class TestRepoRawMeasurements:
    """Test engine against a synthetic raw_measurements CSV."""

    @staticmethod
    def _write_csv(path: Path) -> Path:
        """Write a 5-timestep, 3-dimension CSV matching original repo format."""
        path.write_text(
            "t,x1,x2,x3\n0,0.10,0.20,0.30\n1,0.15,0.25,0.35\n2,0.80,0.75,0.70\n3,0.50,0.55,0.45\n4,0.90,0.85,0.95\n"
        )
        return path

    def test_repo_raw_csv(self, tmp_path: Path) -> None:
        """Process a 5-row raw_measurements CSV end-to-end."""
        csv_path = self._write_csv(tmp_path / "raw_measurements.csv")

        engine = MeasurementEngine(eta=0.10, H_rec=50)
        result = engine.from_csv(csv_path)

        assert result.n_timesteps == 5
        assert result.n_dims == 3

        # Tier-1 identities must hold
        for inv in result.invariants:
            assert abs(inv.F - (1 - inv.omega)) < 1e-12, "F = 1 − ω violated"
            assert abs(inv.IC - math.exp(inv.kappa)) < 1e-12, "IC = exp(κ) violated"
            assert inv.F >= inv.IC - 1e-15, "AM-GM violated"

    def test_repo_raw_csv_casepack(self, tmp_path: Path) -> None:
        """Generate a casepack from a synthetic raw_measurements CSV."""
        csv_path = self._write_csv(tmp_path / "raw_measurements.csv")

        engine = MeasurementEngine(eta=0.10, H_rec=50)
        result = engine.from_csv(csv_path)
        cp_dir = tmp_path / "generated_casepack"
        engine.generate_casepack(result, cp_dir)

        assert (cp_dir / "manifest.json").exists()
        assert (cp_dir / "expected" / "psi.csv").exists()
        assert (cp_dir / "expected" / "invariants.json").exists()

        # Validate invariants.json structure
        doc = json.loads((cp_dir / "expected" / "invariants.json").read_text())
        assert len(doc["rows"]) == 5
        for row in doc["rows"]:
            # F = 1 − ω
            assert abs(row["F"] - (1 - row["omega"])) < 1e-10
