jepa-fx-risk/tests/test_regime.py

"""Tests for scripts/prepare_regime.py — HMM regime detector (jepa-fx-risk#13).

TDD: tests first, implementation follows.
"""

import importlib.util
import os
import shutil
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

_SCRIPT = Path(__file__).parent.parent / "scripts" / "prepare_regime.py"

DATA_DIR = Path(__file__).parent.parent / "data" / "processed"
HOURLY = DATA_DIR / "eurusd_hourly.parquet"
DAILY = DATA_DIR / "eurusd_daily.parquet"


def _import():
    spec = importlib.util.spec_from_file_location("prepare_regime", _SCRIPT)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


@pytest.fixture()
def mod():
    return _import()


# ---------------------------------------------------------------------------
# fit_regime_hmm — pure function (doesn't touch disk)
# ---------------------------------------------------------------------------

def _synthetic_rv(seed=42, size=500):
    """Noisy 3-regime vol series: calm→stressed→crisis→calm interleaved."""
    rng = np.random.default_rng(seed)
    low  = np.abs(rng.normal(0.005, 0.001, size=size // 3))
    mid  = np.abs(rng.normal(0.015, 0.003, size=size // 3))
    high = np.abs(rng.normal(0.04,  0.008, size=size - 2 * (size // 3)))
    return np.concatenate([low, mid, high])


class TestFitRegimeHmm:
    def test_returns_integer_labels(self, mod):
        rv = _synthetic_rv(seed=0)
        labels = mod.fit_regime_hmm(rv, n_states=3, random_state=42)
        assert np.issubdtype(labels.dtype, np.integer), f"dtype={labels.dtype}"
        assert len(labels) == len(rv)

    def test_states_are_0_1_2(self, mod):
        rv = _synthetic_rv(seed=1)
        labels = mod.fit_regime_hmm(rv, n_states=3, random_state=42)
        unique = set(labels.tolist())
        assert unique.issubset({0, 1, 2}), f"unexpected states: {unique}"

    def test_deterministic(self, mod):
        rv = _synthetic_rv(seed=7)
        a = mod.fit_regime_hmm(rv, n_states=3, random_state=42)
        b = mod.fit_regime_hmm(rv, n_states=3, random_state=42)
        assert np.array_equal(a, b), "HMM not deterministic with same random_state"

    def test_sorted_by_vol_asc(self, mod):
        # 3 clearly separated noisy clusters; state 0 should be calm, 2 should be crisis.
        rng = np.random.default_rng(42)
        n = 200
        low  = np.abs(rng.normal(0.005, 0.001, n))
        mid  = np.abs(rng.normal(0.015, 0.003, n))
        high = np.abs(rng.normal(0.05,  0.008, n))
        rv = np.concatenate([low, mid, high])
        labels = mod.fit_regime_hmm(rv, n_states=3, random_state=42)
        # Mean regime label in the high-vol section should exceed mean in the low-vol section.
        assert labels[2*n:].mean() > labels[:n].mean(), \
            "crisis section mean regime label should exceed calm section"
        # The calm section should not be labeled as crisis (2) dominantly.
        calm_modal = int(np.bincount(labels[:n]).argmax())
        assert calm_modal < 2, f"calm section mostly labeled {calm_modal}, expected 0 or 1"

    def test_two_states(self, mod):
        rv = _synthetic_rv(seed=0)
        labels = mod.fit_regime_hmm(rv, n_states=2, random_state=42)
        unique = set(labels.tolist())
        assert unique.issubset({0, 1})




# ---------------------------------------------------------------------------
# prepare_regime_df — reads parquet, fits HMM, returns DataFrame
# ---------------------------------------------------------------------------

class TestPrepareRegimeDf:
    @pytest.mark.skipif(not HOURLY.exists(), reason="hourly parquet not available")
    def test_output_columns(self, mod):
        df = mod.prepare_regime_df(str(HOURLY), freq="hourly")
        assert "datetime" in df.columns
        assert "regime" in df.columns

    @pytest.mark.skipif(not HOURLY.exists(), reason="hourly parquet not available")
    def test_regime_values(self, mod):
        df = mod.prepare_regime_df(str(HOURLY), freq="hourly")
        unique = set(df["regime"].tolist())
        assert unique.issubset({0, 1, 2}), f"unexpected regime values: {unique}"

    @pytest.mark.skipif(not HOURLY.exists(), reason="hourly parquet not available")
    def test_no_nulls(self, mod):
        df = mod.prepare_regime_df(str(HOURLY), freq="hourly")
        assert df["regime"].isna().sum() == 0

    @pytest.mark.skipif(not DAILY.exists(), reason="daily parquet not available")
    def test_daily_fallback(self, mod):
        df = mod.prepare_regime_df(str(DAILY), freq="daily")
        assert "regime" in df.columns
        assert set(df["regime"].tolist()).issubset({0, 1, 2})


# ---------------------------------------------------------------------------
# Integration: check that train.py REGIME SEAM exists and is togglable
# ---------------------------------------------------------------------------

class TestTrainPyRegimeSeam:
    def test_enable_regime_env_var_documented(self):
        train_py = Path(__file__).parent.parent / "train.py"
        content = train_py.read_text()
        assert "JEPA_ENABLE_REGIME" in content, "JEPA_ENABLE_REGIME toggle not found in train.py"

    def test_regime_seam_comment_present(self):
        train_py = Path(__file__).parent.parent / "train.py"
        content = train_py.read_text()
        assert "REGIME" in content and "seam" in content.lower(), \
            "agent-editable regime seam marker not found in train.py"