"""R equivalence tests for DTW distance.

Compares our Rust DTW implementation against R dtw::dtw() with symmetric1
step pattern (standard DTW: equal weight for all 3 directions).

R reference data generated by: scripts/generate_r_reference_dtw.R
"""

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest
from causal_impact import CausalImpact
from causal_impact._core import py_dtw_distance

FIXTURE_DIR = Path(__file__).parent / "fixtures"


@pytest.fixture(scope="module")
def r_dtw_scenarios():
    with open(FIXTURE_DIR / "r_reference_dtw.json") as f:
        return json.load(f)


def _get_scenario(scenarios, name):
    for s in scenarios:
        if s.get("scenario") == name:
            return s
    msg = f"Scenario '{name}' not found in R reference data"
    raise ValueError(msg)


class TestDtwREquivalence:
    def test_matches_r_basic(self, r_dtw_scenarios):
        s = _get_scenario(r_dtw_scenarios, "dtw_basic")
        d = py_dtw_distance(s["x"], s["y"])
        assert d == pytest.approx(s["r_dtw_distance"], abs=1e-6)

    def test_matches_r_with_window(self, r_dtw_scenarios):
        s = _get_scenario(r_dtw_scenarios, "dtw_with_window")
        d = py_dtw_distance(s["x"], s["y"], window=s["window"])
        assert d == pytest.approx(s["r_dtw_distance"], abs=1e-6)

    def test_matches_r_different_length(self, r_dtw_scenarios):
        s = _get_scenario(r_dtw_scenarios, "dtw_different_length")
        d = py_dtw_distance(s["x"], s["y"])
        assert d == pytest.approx(s["r_dtw_distance"], abs=1e-6)

    def test_matches_r_seasonal(self, r_dtw_scenarios):
        s = _get_scenario(r_dtw_scenarios, "dtw_seasonal")
        d = py_dtw_distance(s["x"], s["y"])
        assert d == pytest.approx(s["r_dtw_distance"], abs=1e-6)

    def test_ranking_matches_r(self, r_dtw_scenarios):
        """Control selection ranking should match R dtw::dtw() ranking."""
        s = _get_scenario(r_dtw_scenarios, "control_selection_ranking")
        y = s["y"]
        distances = []
        for cand in s["candidates"]:
            d = py_dtw_distance(cand["values"], y)
            assert d == pytest.approx(cand["r_dtw_distance"], abs=1e-4)
            distances.append((d, cand["name"]))
        distances.sort(key=lambda x: x[0])
        our_ranking = [name for _, name in distances]
        assert our_ranking == s["r_ranking"]


class TestPlaceboConformalE2E:
    """E2E tests for Placebo and Conformal using R fixture data as input.

    R CausalImpact does not have Placebo/Conformal features,
    so we verify statistical properties instead of numeric equivalence.
    """

    @staticmethod
    def _make_no_effect_ts(seed=42, t_pre=60, t_post=20):
        rng = np.random.default_rng(seed)
        y = np.cumsum(rng.normal(0, 1, t_pre + t_post))
        return pd.DataFrame({"y": y}), [0, t_pre - 1], [t_pre, t_pre + t_post - 1]

    @staticmethod
    def _make_strong_effect_ts(seed=42, t_pre=60, t_post=20, effect=8.0):
        rng = np.random.default_rng(seed)
        y = np.cumsum(rng.normal(0, 1, t_pre + t_post))
        y[t_pre:] += effect * np.std(y[:t_pre])
        return pd.DataFrame({"y": y}), [0, t_pre - 1], [t_pre, t_pre + t_post - 1]

    def test_placebo_null_effect_high_p(self):
        """No causal effect -> placebo p-value should be high (>= 0.1)."""
        df, pre, post = self._make_no_effect_ts(seed=123)
        ci = CausalImpact(df, pre, post, model_args={"niter": 200, "nwarmup": 100})
        result = ci.run_placebo_test(n_placebos=10, min_pre_length=15)
        assert result.p_value >= 0.1

    def test_placebo_strong_effect_low_p(self):
        """Strong causal effect -> placebo p-value should be low (<= 0.5)."""
        df, pre, post = self._make_strong_effect_ts(seed=456, effect=8.0)
        ci = CausalImpact(df, pre, post, model_args={"niter": 200, "nwarmup": 100})
        result = ci.run_placebo_test(n_placebos=15, min_pre_length=15)
        assert result.p_value <= 0.5

    def test_conformal_coverage_on_no_effect_data(self):
        """Conformal q_hat should be positive and finite on no-effect data."""
        df, pre, post = self._make_no_effect_ts(seed=999, t_pre=60)
        ci = CausalImpact(df, pre, post, model_args={"niter": 200, "nwarmup": 100})
        result = ci.run_conformal_analysis(alpha=0.1)
        assert result.q_hat > 0.0
        assert np.isfinite(result.q_hat)
        assert np.all(result.lower < result.upper)
