"""
Tests for pre-trends power analysis module.

Tests the implementation of Roth (2022) methods for assessing
the informativeness of pre-trends tests.
"""

import numpy as np
import pandas as pd
import pytest

from diff_diff import MultiPeriodDiD
from diff_diff.pretrends import (
    PreTrendsPower,
    PreTrendsPowerCurve,
    PreTrendsPowerResults,
    compute_mdv,
    compute_pretrends_power,
)
from diff_diff.results import MultiPeriodDiDResults, PeriodEffect

# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def simple_panel_data():
    """Generate simple panel data for testing."""
    np.random.seed(42)
    n_units = 100
    n_periods = 8
    treatment_time = 4
    true_att = 5.0

    data = []
    for unit in range(n_units):
        is_treated = unit < n_units // 2
        unit_effect = np.random.normal(0, 2)

        for period in range(n_periods):
            time_effect = period * 1.0
            y = 10.0 + unit_effect + time_effect

            post = period >= treatment_time
            if is_treated and post:
                y += true_att

            y += np.random.normal(0, 0.5)

            data.append(
                {
                    "unit": unit,
                    "period": period,
                    "treated": int(is_treated),
                    "post": int(post),
                    "outcome": y,
                }
            )

    return pd.DataFrame(data)


@pytest.fixture
def multiperiod_results(simple_panel_data):
    """Fit MultiPeriodDiD and return results."""
    mp_did = MultiPeriodDiD()
    results = mp_did.fit(
        simple_panel_data,
        outcome="outcome",
        treatment="treated",
        time="period",
        post_periods=[4, 5, 6, 7],
    )
    return results


@pytest.fixture
def mock_multiperiod_results():
    """Create mock MultiPeriodDiDResults for unit testing.

    This fixture simulates event study results with:
    - Pre-periods: 0, 1, 2, 3 (period 3 is reference, omitted from estimation)
    - Post-periods: 4, 5, 6, 7
    - Estimated coefficients for periods 0, 1, 2 (pre) and 4, 5, 6, 7 (post)
    """
    # Pre-period effects (excluding reference period 3)
    period_effects = {
        0: PeriodEffect(
            period=0, effect=0.1, se=0.5, t_stat=0.2, p_value=0.84, conf_int=(-0.88, 1.08)
        ),
        1: PeriodEffect(
            period=1, effect=-0.05, se=0.5, t_stat=-0.1, p_value=0.92, conf_int=(-1.03, 0.93)
        ),
        2: PeriodEffect(
            period=2, effect=0.08, se=0.5, t_stat=0.16, p_value=0.87, conf_int=(-0.90, 1.06)
        ),
        # Period 3 is reference - not in period_effects
        # Post-period effects
        4: PeriodEffect(
            period=4, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98)
        ),
        5: PeriodEffect(
            period=5, effect=5.2, se=0.5, t_stat=10.4, p_value=0.0001, conf_int=(4.22, 6.18)
        ),
        6: PeriodEffect(
            period=6, effect=4.8, se=0.5, t_stat=9.6, p_value=0.0001, conf_int=(3.82, 5.78)
        ),
        7: PeriodEffect(
            period=7, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98)
        ),
    }

    # Coefficients for estimated periods (excludes reference period 3)
    coefficients = {
        "treated:period_0": 0.1,
        "treated:period_1": -0.05,
        "treated:period_2": 0.08,
        "treated:period_4": 5.0,
        "treated:period_5": 5.2,
        "treated:period_6": 4.8,
        "treated:period_7": 5.0,
    }

    # Create vcov matrix (diagonal for simplicity)
    # 7 coefficients: 3 pre + 4 post
    vcov = np.diag([0.25] * 7)

    return MultiPeriodDiDResults(
        period_effects=period_effects,
        avg_att=5.0,
        avg_se=0.25,
        avg_t_stat=20.0,
        avg_p_value=0.0001,
        avg_conf_int=(4.51, 5.49),
        n_obs=800,
        n_treated=400,
        n_control=400,
        pre_periods=[0, 1, 2, 3],  # 4 pre-periods, but period 3 is reference
        post_periods=[4, 5, 6, 7],
        vcov=vcov,
        coefficients=coefficients,
    )


# =============================================================================
# Tests for PreTrendsPower class initialization
# =============================================================================


class TestPreTrendsPowerInit:
    """Tests for PreTrendsPower initialization."""

    def test_default_init(self):
        """Test default initialization."""
        pt = PreTrendsPower()
        assert pt.alpha == 0.05
        assert pt.target_power == 0.80
        assert pt.violation_type == "linear"
        assert pt.violation_weights is None

    def test_custom_alpha(self):
        """Test initialization with custom alpha."""
        pt = PreTrendsPower(alpha=0.10)
        assert pt.alpha == 0.10

    def test_custom_power(self):
        """Test initialization with custom target power."""
        pt = PreTrendsPower(power=0.90)
        assert pt.target_power == 0.90

    def test_violation_type_constant(self):
        """Test initialization with constant violation type."""
        pt = PreTrendsPower(violation_type="constant")
        assert pt.violation_type == "constant"

    def test_violation_type_last_period(self):
        """Test initialization with last_period violation type."""
        pt = PreTrendsPower(violation_type="last_period")
        assert pt.violation_type == "last_period"

    def test_violation_type_custom(self):
        """Test initialization with custom violation type."""
        weights = np.array([0.5, 0.3, 0.2])
        pt = PreTrendsPower(violation_type="custom", violation_weights=weights)
        assert pt.violation_type == "custom"
        np.testing.assert_array_equal(pt.violation_weights, weights)

    def test_invalid_alpha_raises(self):
        """Test that invalid alpha raises ValueError."""
        with pytest.raises(ValueError, match="between 0 and 1"):
            PreTrendsPower(alpha=1.5)

    def test_invalid_power_raises(self):
        """Test that invalid power raises ValueError."""
        with pytest.raises(ValueError, match="between 0 and 1"):
            PreTrendsPower(power=0.0)

    def test_invalid_violation_type_raises(self):
        """Test that invalid violation type raises ValueError."""
        with pytest.raises(ValueError, match="violation_type must be"):
            PreTrendsPower(violation_type="invalid")

    def test_custom_without_weights_raises(self):
        """Test that custom type without weights raises ValueError."""
        with pytest.raises(ValueError, match="violation_weights must be provided"):
            PreTrendsPower(violation_type="custom")


# =============================================================================
# Tests for violation weight computation
# =============================================================================


class TestViolationWeights:
    """Tests for violation weight computation."""

    def test_linear_weights(self):
        """Test linear violation weights."""
        pt = PreTrendsPower(violation_type="linear")
        weights = pt._get_violation_weights(4)

        # Should be normalized to unit norm
        assert np.isclose(np.linalg.norm(weights), 1.0)
        # Should increase (in absolute value) as we go back in time
        # For linear: [3, 2, 1, 0] normalized
        assert len(weights) == 4

    def test_constant_weights(self):
        """Constant violation weights are ``[1, 1, ..., 1]`` (no L2 norm).

        REGISTRY ``## PreTrendsPower`` documents ``δ_t = c`` (per-period
        level shift) for the constant violation pattern; PR-B R13 fix
        flipped ``_get_violation_weights('constant')`` to return the
        unnormalized direction so ``δ_t = M`` exactly. The previous
        L2-normalized ``[1/√K, ..., 1/√K]`` direction silently re-scaled
        the reported MDV by ``1/√K`` relative to the documented contract.
        """
        pt = PreTrendsPower(violation_type="constant")
        weights = pt._get_violation_weights(4)

        # PR-B R13: unnormalized [1, 1, 1, 1] (NOT L2-normalized) so
        # δ_t = M reflects a per-period level shift of magnitude M.
        np.testing.assert_allclose(weights, [1.0, 1.0, 1.0, 1.0])
        # L2 norm should be √K, not 1.
        assert np.isclose(np.linalg.norm(weights), 2.0)

    def test_last_period_weights(self):
        """Test last_period violation weights."""
        pt = PreTrendsPower(violation_type="last_period")
        weights = pt._get_violation_weights(4)

        # Only last period should have weight
        assert weights[-1] == 1.0
        assert np.allclose(weights[:-1], 0.0)

    def test_custom_weights(self):
        """Test custom violation weights."""
        custom = np.array([0.6, 0.8])
        pt = PreTrendsPower(violation_type="custom", violation_weights=custom)
        weights = pt._get_violation_weights(2)

        # Should be normalized
        assert np.isclose(np.linalg.norm(weights), 1.0)


# =============================================================================
# Tests for power computation
# =============================================================================


class TestPowerComputation:
    """Tests for power computation."""

    def test_power_at_zero_equals_alpha(self):
        """Test that power at M=0 equals alpha (size of test).

        This is a Wald-form property: under H0, the noncentrality is 0 and
        the rejection probability equals alpha exactly. Under NIS the joint
        rejection probability at H0 is 1 - (1 - alpha)^K ≈ K*alpha for
        small alpha (~0.14 for K=3 at alpha=0.05). Pin Wald to test the
        Wald-specific size property.
        """
        pt = PreTrendsPower(alpha=0.05, pretest_form="wald")

        # Create simple vcov
        n_pre = 3
        vcov = np.eye(n_pre) * 0.25
        weights = pt._get_violation_weights(n_pre)

        power, _, _, _ = pt._compute_power(0.0, weights, vcov)
        assert np.isclose(power, 0.05, atol=0.01)

    def test_power_increases_with_M(self):
        """Test that power increases with M."""
        pt = PreTrendsPower()

        n_pre = 3
        vcov = np.eye(n_pre) * 0.25
        weights = pt._get_violation_weights(n_pre)

        power_small, _, _, _ = pt._compute_power(0.5, weights, vcov)
        power_large, _, _, _ = pt._compute_power(2.0, weights, vcov)

        assert power_large > power_small

    def test_power_at_mdv_equals_target(self):
        """Test that power at MDV equals target power."""
        pt = PreTrendsPower(power=0.80)

        n_pre = 3
        vcov = np.eye(n_pre) * 0.25
        weights = pt._get_violation_weights(n_pre)

        mdv = pt._compute_mdv(weights, vcov)
        power, _, _, _ = pt._compute_power(mdv, weights, vcov)

        assert np.isclose(power, 0.80, atol=0.02)


# =============================================================================
# Tests for MDV computation
# =============================================================================


class TestMDVComputation:
    """Tests for minimum detectable violation computation."""

    def test_mdv_positive(self):
        """Test that MDV is positive."""
        pt = PreTrendsPower()

        n_pre = 3
        vcov = np.eye(n_pre) * 0.25
        weights = pt._get_violation_weights(n_pre)

        mdv = pt._compute_mdv(weights, vcov)
        assert mdv > 0

    def test_mdv_decreases_with_precision(self):
        """Test that MDV decreases with smaller variance."""
        pt = PreTrendsPower()

        n_pre = 3
        weights = pt._get_violation_weights(n_pre)

        # Large variance
        vcov_large = np.eye(n_pre) * 1.0
        mdv_large = pt._compute_mdv(weights, vcov_large)

        # Small variance
        vcov_small = np.eye(n_pre) * 0.1
        mdv_small = pt._compute_mdv(weights, vcov_small)

        assert mdv_small < mdv_large

    def test_mdv_increases_with_target_power(self):
        """Test that MDV increases with higher target power."""
        n_pre = 3
        vcov = np.eye(n_pre) * 0.25

        pt_low = PreTrendsPower(power=0.50)
        pt_high = PreTrendsPower(power=0.90)

        weights = pt_low._get_violation_weights(n_pre)

        mdv_low = pt_low._compute_mdv(weights, vcov)
        mdv_high = pt_high._compute_mdv(weights, vcov)

        assert mdv_high > mdv_low


# =============================================================================
# Tests for fit method
# =============================================================================


class TestPreTrendsPowerFit:
    """Tests for PreTrendsPower.fit() method."""

    def test_fit_returns_results(self, mock_multiperiod_results):
        """Test that fit returns PreTrendsPowerResults."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        assert isinstance(results, PreTrendsPowerResults)

    def test_fit_custom_M(self, mock_multiperiod_results):
        """Test fit with custom M value."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results, M=1.0)

        assert results.violation_magnitude == 1.0

    def test_results_has_expected_attributes(self, mock_multiperiod_results):
        """Test that results have all expected attributes."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        assert hasattr(results, "power")
        assert hasattr(results, "mdv")
        assert hasattr(results, "violation_magnitude")
        assert hasattr(results, "violation_type")
        assert hasattr(results, "alpha")
        assert hasattr(results, "target_power")
        assert hasattr(results, "n_pre_periods")
        assert hasattr(results, "test_statistic")
        assert hasattr(results, "critical_value")
        assert hasattr(results, "noncentrality")

    def test_results_n_pre_periods(self, mock_multiperiod_results):
        """Test that n_pre_periods matches estimated pre-periods (excluding reference)."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        # n_pre_periods should be the number of estimated coefficients (3)
        # not the total number of pre-periods (4), since period 3 is the reference
        expected_n_pre = len(
            [
                p
                for p in mock_multiperiod_results.pre_periods
                if f"treated:period_{p}" in mock_multiperiod_results.coefficients
            ]
        )
        assert results.n_pre_periods == expected_n_pre
        assert results.n_pre_periods == 3  # 4 pre-periods minus 1 reference


# =============================================================================
# Tests for power_curve method
# =============================================================================


class TestPowerCurve:
    """Tests for power_curve method."""

    def test_power_curve_returns_curve(self, mock_multiperiod_results):
        """Test that power_curve returns PreTrendsPowerCurve."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results)

        assert isinstance(curve, PreTrendsPowerCurve)

    def test_power_curve_custom_grid(self, mock_multiperiod_results):
        """Test power_curve with custom M grid."""
        pt = PreTrendsPower()
        M_grid = [0, 0.5, 1.0, 1.5, 2.0]
        curve = pt.power_curve(mock_multiperiod_results, M_grid=M_grid)

        assert len(curve.M_values) == len(M_grid)
        np.testing.assert_array_equal(curve.M_values, M_grid)

    def test_power_curve_n_points(self, mock_multiperiod_results):
        """Test power_curve with custom n_points."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results, n_points=20)

        assert len(curve.M_values) == 20

    def test_power_curve_monotonic(self, mock_multiperiod_results):
        """Test that power curve is monotonically increasing."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results)

        # Power should be non-decreasing
        diffs = np.diff(curve.powers)
        assert np.all(diffs >= -0.01)  # Allow small numerical noise

    def test_power_curve_to_dataframe(self, mock_multiperiod_results):
        """Test power_curve to_dataframe method."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results)
        df = curve.to_dataframe()

        assert isinstance(df, pd.DataFrame)
        assert "M" in df.columns
        assert "power" in df.columns


# =============================================================================
# Tests for PreTrendsPowerResults
# =============================================================================


class TestPreTrendsPowerResults:
    """Tests for PreTrendsPowerResults dataclass."""

    def test_results_summary(self, mock_multiperiod_results):
        """Test summary method produces string."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        summary = results.summary()
        assert isinstance(summary, str)
        assert "Pre-Trends Power Analysis" in summary
        assert "Roth 2022" in summary

    def test_results_to_dict(self, mock_multiperiod_results):
        """Test to_dict method."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        d = results.to_dict()
        assert isinstance(d, dict)
        assert "power" in d
        assert "mdv" in d
        assert "violation_type" in d

    def test_results_to_dataframe(self, mock_multiperiod_results):
        """Test to_dataframe method."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        df = results.to_dataframe()
        assert isinstance(df, pd.DataFrame)
        assert len(df) == 1

    def test_is_informative_property(self, mock_multiperiod_results):
        """Test is_informative property."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        assert isinstance(results.is_informative, bool)

    def test_power_adequate_property(self, mock_multiperiod_results):
        """Test power_adequate property."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        assert isinstance(results.power_adequate, bool)

    def test_power_at_works_for_custom_violation_type(self, mock_multiperiod_results):
        """power_at(M) now works for custom violation type (PR-B Step 5).

        PR-A R18 added a NotImplementedError guard because
        PreTrendsPowerResults did not persist fitted violation_weights.
        PR-B persisted them on the result dataclass and refactored
        power_at() to read them directly. This test confirms the guard
        is lifted for fresh fits: a custom-weights PreTrendsPower fit
        produces a result whose power_at(M) returns a finite, in-[0,1]
        power value.
        """
        weights = np.array([0.1, 0.3, 0.6])
        pt = PreTrendsPower(violation_type="custom", violation_weights=weights)
        results = pt.fit(mock_multiperiod_results)

        # No longer raises; returns a finite power value in [0, 1].
        power = results.power_at(0.5)
        assert np.isfinite(power)
        assert 0.0 <= power <= 1.0

    def test_power_at_raises_on_legacy_custom_result_without_weights(
        self, mock_multiperiod_results
    ):
        """power_at(M) still raises for old serialized results lacking
        violation_weights (backwards-compat guard).

        The dataclass default for violation_weights is None; old serialized
        PreTrendsPowerResults objects from before PR-B's field addition will
        have None there. For custom fits, power_at() cannot reconstruct
        custom weights from violation_type + n_pre_periods alone, so the
        PR-A R18 guard is retained for that specific backwards-compat path.
        """
        weights = np.array([0.1, 0.3, 0.6])
        pt = PreTrendsPower(violation_type="custom", violation_weights=weights)
        results = pt.fit(mock_multiperiod_results)
        # Simulate a legacy-result scenario by clearing the persisted weights.
        results.violation_weights = None

        with pytest.raises(NotImplementedError, match="custom violation weights"):
            results.power_at(0.5)


# =============================================================================
# Tests for convenience functions
# =============================================================================


class TestConvenienceFunctions:
    """Tests for convenience functions."""

    def test_compute_pretrends_power(self, mock_multiperiod_results):
        """Test compute_pretrends_power function."""
        results = compute_pretrends_power(mock_multiperiod_results)

        assert isinstance(results, PreTrendsPowerResults)

    def test_compute_pretrends_power_custom_params(self, mock_multiperiod_results):
        """Test compute_pretrends_power with custom parameters."""
        results = compute_pretrends_power(
            mock_multiperiod_results, alpha=0.10, target_power=0.90, violation_type="constant"
        )

        assert results.alpha == 0.10
        assert results.target_power == 0.90
        assert results.violation_type == "constant"

    def test_compute_mdv(self, mock_multiperiod_results):
        """Test compute_mdv function."""
        mdv = compute_mdv(mock_multiperiod_results)

        assert isinstance(mdv, float)
        assert mdv > 0

    def test_compute_pretrends_power_rejects_custom_violation_type(self, mock_multiperiod_results):
        """compute_pretrends_power(..., violation_type='custom') without explicit
        ``violation_weights`` must raise ValueError.

        PR-B Step 6 added the ``violation_weights`` kwarg to both helpers, so
        ``violation_type='custom'`` is now usable from the helper API when the
        weights vector is supplied. This regression locks the loud-fail
        contract for the unsupplied-weights case: silently coercing a custom
        request into a degenerate (zero / equal-weights) fit was the PR-A
        R18 silent-failure that the loud guard prevents. See REGISTRY.md
        PreTrendsPower section + docs/methodology/papers/roth-2022-review.md.
        """
        with pytest.raises(ValueError, match="violation_weights"):
            compute_pretrends_power(mock_multiperiod_results, violation_type="custom")

    def test_compute_mdv_rejects_custom_violation_type(self, mock_multiperiod_results):
        """compute_mdv(..., violation_type='custom') without ``violation_weights``
        must raise ValueError.

        Same contract as ``compute_pretrends_power``: PR-B Step 6 made the
        helper accept ``violation_weights``, so the rejection is now scoped
        to the unsupplied-weights case rather than the entire custom path.
        """
        with pytest.raises(ValueError, match="violation_weights"):
            compute_mdv(mock_multiperiod_results, violation_type="custom")


# =============================================================================
# Tests for get_params and set_params
# =============================================================================


class TestGetSetParams:
    """Tests for sklearn-like parameter interface."""

    def test_get_params(self):
        """Test get_params method."""
        pt = PreTrendsPower(alpha=0.10, power=0.90, violation_type="constant")
        params = pt.get_params()

        assert params["alpha"] == 0.10
        assert params["power"] == 0.90
        assert params["violation_type"] == "constant"

    def test_set_params(self):
        """Test set_params method."""
        pt = PreTrendsPower()
        pt.set_params(alpha=0.10, power=0.90)

        assert pt.alpha == 0.10
        assert pt.target_power == 0.90


# =============================================================================
# Tests for integration with real estimators
# =============================================================================


class TestIntegration:
    """Integration tests with event study results.

    Note: These tests use mock results with pre-period coefficients.
    MultiPeriodDiD by default only estimates post-period treatment effects,
    so we use a fixture that simulates full event study results with
    pre-period coefficients (excluding the reference period).
    """

    def test_with_multiperiod_results(self, mock_multiperiod_results):
        """Test full pipeline with MultiPeriodDiDResults."""
        # Run pre-trends power analysis
        pt = PreTrendsPower()
        power_results = pt.fit(mock_multiperiod_results)

        # Check results are reasonable
        assert power_results.power >= 0
        assert power_results.power <= 1
        assert power_results.mdv > 0 or np.isinf(power_results.mdv)

    def test_power_curve_with_results(self, mock_multiperiod_results):
        """Test power curve with event study results."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results, n_points=10)

        # Check curve properties
        assert len(curve.M_values) == 10
        assert len(curve.powers) == 10
        assert np.all((curve.powers >= 0) & (curve.powers <= 1))

    def test_sensitivity_to_honest_did(self, mock_multiperiod_results):
        """Test sensitivity_to_honest_did method."""
        pt = PreTrendsPower()
        sensitivity = pt.sensitivity_to_honest_did(mock_multiperiod_results)

        assert "mdv" in sensitivity
        assert "interpretation" in sensitivity
        assert isinstance(sensitivity["interpretation"], str)


# =============================================================================
# Tests for different violation types
# =============================================================================


class TestViolationTypes:
    """Tests for different violation types."""

    def test_linear_violation(self, mock_multiperiod_results):
        """Test power analysis with linear violation."""
        pt = PreTrendsPower(violation_type="linear")
        results = pt.fit(mock_multiperiod_results)

        assert results.violation_type == "linear"

    def test_constant_violation(self, mock_multiperiod_results):
        """Test power analysis with constant violation."""
        pt = PreTrendsPower(violation_type="constant")
        results = pt.fit(mock_multiperiod_results)

        assert results.violation_type == "constant"

    def test_last_period_violation(self, mock_multiperiod_results):
        """Test power analysis with last_period violation."""
        pt = PreTrendsPower(violation_type="last_period")
        results = pt.fit(mock_multiperiod_results)

        assert results.violation_type == "last_period"

    def test_different_types_give_different_results(self, mock_multiperiod_results):
        """Test that different violation types can give different MDV."""
        pt_linear = PreTrendsPower(violation_type="linear")
        pt_constant = PreTrendsPower(violation_type="constant")
        pt_last = PreTrendsPower(violation_type="last_period")

        mdv_linear = pt_linear.fit(mock_multiperiod_results).mdv
        mdv_constant = pt_constant.fit(mock_multiperiod_results).mdv
        mdv_last = pt_last.fit(mock_multiperiod_results).mdv

        # All MDVs should be positive (or inf for degenerate cases)
        assert mdv_linear > 0 or np.isinf(mdv_linear)
        assert mdv_constant > 0 or np.isinf(mdv_constant)
        assert mdv_last > 0 or np.isinf(mdv_last)

        # Just verify we can compute all three types without error
        # (with diagonal vcov, they may be the same due to normalization)


# =============================================================================
# Tests for edge cases
# =============================================================================


class TestEdgeCases:
    """Tests for edge cases."""

    def test_single_pre_period(self):
        """Test with single pre-period (excluding reference).

        This tests the case where there's only one estimated pre-period
        coefficient. We have pre_periods=[2, 3] where period 3 is the
        reference (excluded), leaving only period 2 as estimated.
        """
        period_effects = {
            2: PeriodEffect(
                period=2, effect=0.1, se=0.5, t_stat=0.2, p_value=0.84, conf_int=(-0.88, 1.08)
            ),
            # Period 3 is reference - not estimated
            4: PeriodEffect(
                period=4, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98)
            ),
        }

        coefficients = {
            "treated:period_2": 0.1,
            "treated:period_4": 5.0,
        }

        results = MultiPeriodDiDResults(
            period_effects=period_effects,
            avg_att=5.0,
            avg_se=0.5,
            avg_t_stat=10.0,
            avg_p_value=0.0001,
            avg_conf_int=(4.02, 5.98),
            n_obs=200,
            n_treated=100,
            n_control=100,
            pre_periods=[2, 3],  # Period 3 is reference
            post_periods=[4],
            vcov=np.array([[0.25, 0], [0, 0.25]]),
            coefficients=coefficients,
        )

        pt = PreTrendsPower()
        power_results = pt.fit(results)

        assert power_results.n_pre_periods == 1  # Only period 2 is estimated

    def test_many_pre_periods(self):
        """Test with many pre-periods.

        This tests the case with 10 pre-periods where period 9 is the
        reference (excluded), leaving 9 estimated pre-period coefficients.
        """
        n_pre_total = 10
        n_pre_estimated = 9  # Excluding reference period

        # Pre-period effects (excluding reference period 9)
        period_effects = {}
        for i in range(n_pre_estimated):
            period_effects[i] = PeriodEffect(
                period=i,
                effect=0.05 * (i - 4),
                se=0.5,
                t_stat=0.1 * (i - 4),
                p_value=0.92,
                conf_int=(-0.88, 1.08),
            )

        # Post-period effects
        for i in range(4):
            period_effects[n_pre_total + i] = PeriodEffect(
                period=n_pre_total + i,
                effect=5.0,
                se=0.5,
                t_stat=10.0,
                p_value=0.0001,
                conf_int=(4.02, 5.98),
            )

        # Coefficients (excluding reference period 9)
        coefficients = {}
        for i in range(n_pre_estimated):
            coefficients[f"treated:period_{i}"] = 0.05 * (i - 4)
        for i in range(4):
            coefficients[f"treated:period_{n_pre_total + i}"] = 5.0

        results = MultiPeriodDiDResults(
            period_effects=period_effects,
            avg_att=5.0,
            avg_se=0.5,
            avg_t_stat=10.0,
            avg_p_value=0.0001,
            avg_conf_int=(4.02, 5.98),
            n_obs=200,
            n_treated=100,
            n_control=100,
            pre_periods=list(range(n_pre_total)),  # Includes reference period 9
            post_periods=list(range(n_pre_total, n_pre_total + 4)),
            vcov=np.diag([0.25] * (n_pre_estimated + 4)),
            coefficients=coefficients,
        )

        pt = PreTrendsPower()
        power_results = pt.fit(results)

        assert power_results.n_pre_periods == n_pre_estimated

    def test_unsupported_results_type_raises(self):
        """Test that unsupported results type raises TypeError."""
        pt = PreTrendsPower()

        with pytest.raises(TypeError, match="Unsupported results type"):
            pt.fit("not a results object")

    def test_callaway_santanna_universal_base_period(self):
        """Test that reference period (e=-1) is correctly filtered out with universal base period.

        The reference period has n_groups=0 and se=NaN, so it should be excluded
        from pre-trends power analysis to avoid contaminating the vcov matrix.
        """
        from diff_diff import CallawaySantAnna, generate_staggered_data

        # Generate data and fit with universal base period
        data = generate_staggered_data(n_units=200, n_periods=10, seed=42)
        cs = CallawaySantAnna(base_period="universal")
        results = cs.fit(
            data,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
        )

        # Verify reference period exists with NaN SE
        assert -1 in results.event_study_effects
        assert np.isnan(results.event_study_effects[-1]["se"])

        # PreTrendsPower should work without errors (reference period filtered out)
        pt = PreTrendsPower()
        power_results = pt.fit(results)

        # Should have valid (non-NaN) results
        assert np.isfinite(power_results.power)
        assert power_results.n_pre_periods >= 1


# =============================================================================
# Tests for visualization (without rendering)
# =============================================================================


class TestVisualization:
    """Tests for visualization methods (without rendering)."""

    def test_power_curve_has_plot_method(self, mock_multiperiod_results):
        """Test that PreTrendsPowerCurve has plot method."""
        pt = PreTrendsPower()
        curve = pt.power_curve(mock_multiperiod_results)

        assert hasattr(curve, "plot")
        assert callable(curve.plot)


# =============================================================================
# Tests for PreTrendsPowerResults.power_at() method
# =============================================================================


class TestPreTrendsPowerResultsPowerAt:
    """Tests for the power_at method on PreTrendsPowerResults."""

    def test_power_at_basic(self, mock_multiperiod_results):
        """Test basic power_at functionality."""
        pt = PreTrendsPower()
        results = pt.fit(mock_multiperiod_results)

        # Compute power at different M values
        power_1 = results.power_at(1.0)
        power_2 = results.power_at(2.0)
        power_5 = results.power_at(5.0)

        # Power should increase with M
        assert power_1 < power_2 < power_5

        # Power should be between 0 and 1
        assert 0 <= power_1 <= 1
        assert 0 <= power_2 <= 1
        assert 0 <= power_5 <= 1

    def test_power_at_zero(self, mock_multiperiod_results):
        """Test power_at with M=0 (should equal alpha under Wald form).

        See note on TestPowerComputation.test_power_at_zero_equals_alpha:
        the exact-equals-alpha property is Wald-specific. Pin Wald.
        """
        pt = PreTrendsPower(alpha=0.05, pretest_form="wald")
        results = pt.fit(mock_multiperiod_results)

        power_0 = results.power_at(0.0)

        # At M=0, power should equal size (alpha) under Wald.
        assert np.isclose(power_0, 0.05, atol=0.01)

    def test_power_at_matches_fit(self, mock_multiperiod_results):
        """Test that power_at gives same result as fitting with that M."""
        pt = PreTrendsPower()

        # Get results from fit
        results1 = pt.fit(mock_multiperiod_results, M=2.0)

        # Get power from power_at method
        results_base = pt.fit(mock_multiperiod_results)
        power_from_method = results_base.power_at(2.0)

        # Should be the same (or very close)
        assert np.isclose(results1.power, power_from_method, rtol=0.01)

    def test_power_at_linear_weights(self, mock_multiperiod_results):
        """Test power_at uses correct linear weights."""
        pt = PreTrendsPower(violation_type="linear")
        results = pt.fit(mock_multiperiod_results)

        # Power_at should work without error
        power = results.power_at(1.0)
        assert 0 <= power <= 1

    def test_power_at_constant_weights(self, mock_multiperiod_results):
        """Test power_at uses correct constant weights."""
        pt = PreTrendsPower(violation_type="constant")
        results = pt.fit(mock_multiperiod_results)

        power = results.power_at(1.0)
        assert 0 <= power <= 1

    def test_power_at_last_period_weights(self, mock_multiperiod_results):
        """Test power_at uses correct last_period weights."""
        pt = PreTrendsPower(violation_type="last_period")
        results = pt.fit(mock_multiperiod_results)

        power = results.power_at(1.0)
        assert 0 <= power <= 1


# =============================================================================
# Tests for pre_periods parameter
# =============================================================================


class TestPrePeriodsParameter:
    """Tests for the pre_periods parameter in fit and related methods."""

    @pytest.fixture
    def event_study_all_periods_results(self):
        """Create results simulating all periods estimated as post_periods.

        This mimics the event study workflow where we estimate coefficients
        for ALL periods (pre and post) to get pre-period placebo effects.
        """
        # Periods 0-3 are pre-treatment, 4-7 are post
        # But we estimate ALL periods as "post" to get coefficients
        period_effects = {}
        coefficients = {}

        # Pre-periods (0, 1, 2) - period 3 would be reference
        for p in [0, 1, 2]:
            period_effects[p] = PeriodEffect(
                period=p,
                effect=np.random.normal(0, 0.1),
                se=0.5,
                t_stat=0.2,
                p_value=0.84,
                conf_int=(-0.88, 1.08),
            )
            coefficients[f"treated:period_{p}"] = period_effects[p].effect

        # Post-periods (4, 5, 6, 7)
        for p in [4, 5, 6, 7]:
            period_effects[p] = PeriodEffect(
                period=p,
                effect=5.0 + np.random.normal(0, 0.1),
                se=0.5,
                t_stat=10.0,
                p_value=0.0001,
                conf_int=(4.02, 5.98),
            )
            coefficients[f"treated:period_{p}"] = period_effects[p].effect

        # In this scenario, pre_periods=[3] (only reference), post_periods=[0,1,2,4,5,6,7]
        vcov = np.diag([0.25] * 7)

        return MultiPeriodDiDResults(
            period_effects=period_effects,
            avg_att=5.0,
            avg_se=0.25,
            avg_t_stat=20.0,
            avg_p_value=0.0001,
            avg_conf_int=(4.51, 5.49),
            n_obs=800,
            n_treated=400,
            n_control=400,
            pre_periods=[3],  # Only reference period
            post_periods=[0, 1, 2, 4, 5, 6, 7],  # All estimated periods
            vcov=vcov,
            coefficients=coefficients,
        )

    def test_fit_with_explicit_pre_periods(self, event_study_all_periods_results):
        """Test fit() with explicit pre_periods parameter."""
        pt = PreTrendsPower()

        # Without pre_periods, would fail because results.pre_periods=[3]
        # and period 3 has no coefficient (it's the reference)
        # With explicit pre_periods=[0,1,2], should work
        results = pt.fit(event_study_all_periods_results, pre_periods=[0, 1, 2])

        assert results.n_pre_periods == 3
        assert results.power >= 0
        assert results.mdv > 0

    def test_pre_periods_overrides_results(self, event_study_all_periods_results):
        """Test that pre_periods parameter overrides results.pre_periods."""
        pt = PreTrendsPower()

        # Explicitly set pre_periods to [0, 1]
        results = pt.fit(event_study_all_periods_results, pre_periods=[0, 1])

        # Should use 2 pre-periods, not what's in results
        assert results.n_pre_periods == 2

    def test_power_at_with_pre_periods(self, event_study_all_periods_results):
        """Test power_at() method with pre_periods parameter."""
        pt = PreTrendsPower()

        power = pt.power_at(event_study_all_periods_results, M=1.0, pre_periods=[0, 1, 2])

        assert 0 <= power <= 1

    def test_power_curve_with_pre_periods(self, event_study_all_periods_results):
        """Test power_curve() with pre_periods parameter."""
        pt = PreTrendsPower()

        curve = pt.power_curve(event_study_all_periods_results, n_points=10, pre_periods=[0, 1, 2])

        assert len(curve.M_values) == 10
        assert len(curve.powers) == 10

    def test_sensitivity_to_honest_did_with_pre_periods(self, event_study_all_periods_results):
        """Test sensitivity_to_honest_did() with pre_periods parameter."""
        pt = PreTrendsPower()

        sensitivity = pt.sensitivity_to_honest_did(
            event_study_all_periods_results, pre_periods=[0, 1, 2]
        )

        assert "mdv" in sensitivity
        assert sensitivity["mdv"] > 0

    def test_convenience_functions_with_pre_periods(self, event_study_all_periods_results):
        """Test convenience functions with pre_periods parameter."""
        # compute_mdv
        mdv = compute_mdv(event_study_all_periods_results, pre_periods=[0, 1, 2])
        assert mdv > 0

        # compute_pretrends_power
        results = compute_pretrends_power(
            event_study_all_periods_results, M=1.0, pre_periods=[0, 1, 2]
        )
        assert results.n_pre_periods == 3
