"""Tests for Markov regime detection - Rust HMM + Python pipeline wrapper."""

import math
import pytest
from horizon._horizon import MarkovRegimeModel, prices_to_returns


class TestMarkovRegimeModel:
    """Test the Rust HMM directly."""

    def test_create_model(self):
        model = MarkovRegimeModel(n_states=2)
        assert model.n_states() == 2
        assert not model.is_trained()

    def test_create_three_state(self):
        model = MarkovRegimeModel(n_states=3)
        assert model.n_states() == 3

    def test_invalid_states(self):
        with pytest.raises(ValueError):
            MarkovRegimeModel(n_states=1)
        with pytest.raises(ValueError):
            MarkovRegimeModel(n_states=9)

    def test_fit_basic(self):
        # Two regimes: low vol then high vol
        data = []
        for i in range(200):
            noise = ((i * 1103515245 + 12345) % 1000 - 500) / 10000.0
            sigma = 0.01 if i < 100 else 0.05
            data.append(noise * sigma)

        model = MarkovRegimeModel(n_states=2)
        ll = model.fit(data, max_iters=50)
        assert math.isfinite(ll)
        assert model.is_trained()

    def test_fit_requires_min_data(self):
        model = MarkovRegimeModel(n_states=2)
        with pytest.raises(ValueError):
            model.fit([0.01])

    def test_fit_rejects_nan(self):
        model = MarkovRegimeModel(n_states=2)
        with pytest.raises(ValueError):
            model.fit([0.01, float("nan"), 0.02])

    def test_emission_params_sorted_by_variance(self):
        """State 0 should have lowest variance after training."""
        data = []
        for i in range(200):
            noise = ((i * 6364136223846793005 + 1) % 1000 - 500) / 10000.0
            sigma = 0.005 if i < 100 else 0.05
            data.append(noise * sigma)

        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=50)
        params = model.emission_params()
        assert params[0][1] <= params[1][1]  # var[0] <= var[1]

    def test_decode_returns_states(self):
        data = [0.001 * ((-1) ** i) for i in range(100)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        states = model.decode(data)
        assert len(states) == 100
        assert all(s in (0, 1) for s in states)

    def test_decode_before_fit_raises(self):
        model = MarkovRegimeModel(n_states=2)
        with pytest.raises(RuntimeError):
            model.decode([0.01, 0.02])

    def test_filter_step(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)

        probs = model.filter_step(0.001)
        assert len(probs) == 2
        assert abs(sum(probs) - 1.0) < 1e-10

    def test_filter_step_before_fit_raises(self):
        model = MarkovRegimeModel(n_states=2)
        with pytest.raises(RuntimeError):
            model.filter_step(0.01)

    def test_filter_step_rejects_nan(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        with pytest.raises(ValueError):
            model.filter_step(float("nan"))

    def test_predict(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        pred = model.predict()
        assert len(pred) == 2
        assert abs(sum(pred) - 1.0) < 1e-10

    def test_smooth(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        smoothed = model.smooth(data)
        assert len(smoothed) == 50
        for row in smoothed:
            assert abs(sum(row) - 1.0) < 1e-10

    def test_transition_matrix(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        tm = model.transition_matrix()
        assert len(tm) == 2
        for row in tm:
            assert abs(sum(row) - 1.0) < 1e-10

    def test_current_regime(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        regime = model.current_regime()
        assert regime in (0, 1)

    def test_reset_filter(self):
        data = [0.001 * ((-1) ** i) for i in range(50)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)
        # Filter some observations
        for x in data[:10]:
            model.filter_step(x)
        # Reset
        model.reset_filter()
        probs = model.filtered_probs()
        # After reset, should be back to initial (roughly uniform for 2-state)
        assert len(probs) == 2

    def test_repr(self):
        model = MarkovRegimeModel(n_states=2)
        r = repr(model)
        assert "MarkovRegimeModel" in r
        assert "n_states=2" in r


class TestPricesToReturns:
    def test_basic(self):
        prices = [100.0, 101.0, 99.0, 102.0]
        rets = prices_to_returns(prices)
        assert len(rets) == 3
        assert abs(rets[0] - math.log(101.0 / 100.0)) < 1e-10

    def test_too_short(self):
        with pytest.raises(ValueError):
            prices_to_returns([100.0])

    def test_zero_price_handled(self):
        rets = prices_to_returns([0.0, 100.0, 200.0])
        assert len(rets) == 2
        assert rets[0] == 0.0  # zero price -> 0 return


class TestMarkovPipelineWrapper:
    """Test the Python markov_regime() pipeline function."""

    def test_import(self):
        from horizon.markov import markov_regime
        assert callable(markov_regime)

    def test_with_pretrained_model(self):
        from horizon.markov import markov_regime
        from horizon.context import Context, FeedData

        # Train a model
        data = [0.001 * ((-1) ** i) for i in range(100)]
        model = MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)

        fn = markov_regime(model=model, feed="test")

        # Simulate a few ticks
        for i in range(5):
            ctx = Context(
                feeds={"test": FeedData(price=0.50 + 0.001 * i)},
                params={},
            )
            fn(ctx)

        # After enough ticks, regime should be injected
        ctx = Context(
            feeds={"test": FeedData(price=0.505)},
            params={},
        )
        fn(ctx)
        assert "regime" in ctx.params or len(ctx.params) == 0  # may need warmup

    def test_auto_train_mode(self):
        from horizon.markov import markov_regime
        from horizon.context import Context, FeedData

        fn = markov_regime(n_states=2, warmup=20, feed="test")

        # Feed 30 ticks to trigger auto-training
        for i in range(30):
            ctx = Context(
                feeds={"test": FeedData(price=0.50 + 0.001 * (i % 5))},
                params={},
            )
            fn(ctx)

        # After warmup + training, regime should be present
        ctx = Context(
            feeds={"test": FeedData(price=0.505)},
            params={},
        )
        fn(ctx)
        # With auto-train, after 30 ticks with warmup=20, should be trained
        if "regime" in ctx.params:
            assert ctx.params["regime"] in (0, 1)

    def test_hz_backtest_integration(self):
        """Full integration: markov_regime in a backtest pipeline."""
        import horizon as hz

        data = [0.001 * ((-1) ** i) for i in range(100)]
        model = hz.MarkovRegimeModel(n_states=2)
        model.fit(data, max_iters=30)

        def quoter(ctx):
            regime = ctx.params.get("regime", 0)
            spread = 0.06 if regime == 1 else 0.04  # wider in volatile
            return hz.Quote(0.48, 0.48 + spread, 5)

        tick_data = [
            {"timestamp": float(i), "price": 0.50 + 0.001 * (i % 10 - 5)}
            for i in range(50)
        ]

        result = hz.backtest(
            data=tick_data,
            pipeline=[hz.markov_regime(model=model), quoter],
        )
        assert len(result.equity_curve) == 50
