from __future__ import annotations

from pathlib import Path
import sys

import numpy as np

HERE = Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

from common_models import (
    UNIVERSE_SIZES,
    acceptable_indices,
    causal_ewma_adjust,
    fit_models,
    model_quantile_es,
    observation_log_losses,
)
from run_study1 import Config, DGP_NAMES, draw_path, evaluate_replication


def test_nested_universes_and_models() -> None:
    assert UNIVERSE_SIZES == {"U3": 3, "U5": 5, "U9": 9}
    rng = np.random.default_rng(11)
    values = rng.normal(size=1000)
    models = fit_models(values, 12)
    assert len(models) == 9
    losses = observation_log_losses(models, rng.normal(size=120))
    assert losses.shape == (120, 9)
    assert np.isfinite(losses).all()


def test_acceptable_set_contains_winner() -> None:
    losses = np.array(
        [[1.0, 1.1, 2.0], [1.0, 1.2, 2.2], [1.1, 1.0, 2.1]]
    )
    accepted, winner, _, _ = acceptable_indices(losses)
    assert winner in accepted
    assert len(accepted) >= 1


def test_quantile_ordering() -> None:
    rng = np.random.default_rng(22)
    values = rng.normal(size=1000)
    gaussian = fit_models(values, 23)[0]
    quantile, expected_shortfall = model_quantile_es(gaussian, values)
    assert expected_shortfall <= quantile < 0


def test_ewma_is_prefix_invariant_after_initialization() -> None:
    rng = np.random.default_rng(24)
    values = rng.normal(size=1400)
    full_adjusted, full_scales = causal_ewma_adjust(values)
    for cutoff in (1120, 1250, 1400):
        prefix_adjusted, prefix_scales = causal_ewma_adjust(values[:cutoff])
        np.testing.assert_array_equal(
            prefix_adjusted, full_adjusted[:cutoff]
        )
        np.testing.assert_array_equal(prefix_scales, full_scales[:cutoff])


def test_all_dgps_and_cells() -> None:
    config = Config(seed=31, reps_per_dgp=1)
    for dgp in DGP_NAMES:
        path = draw_path(dgp, np.random.default_rng(32), config)
        assert len(path) == 2250
        rows = evaluate_replication((dgp, 0, config))
        assert len(rows) == 15
        assert not any(row["fit_failed"] for row in rows)


if __name__ == "__main__":
    import pytest

    raise SystemExit(pytest.main([__file__, "-q"]))
