"""Test all Python code snippets from docs/markov.mdx.

Each snippet is wrapped in a try/except that prints PASS/FAIL.
Snippets requiring network access, real exchanges, or hz.run() are skipped.
Synthetic two-regime data is used where historical data is needed.
"""

import math
import random
import traceback

import horizon as hz

# --------------------------------------------------------------------------- #
# Generate synthetic two-regime data used by multiple snippets
# --------------------------------------------------------------------------- #
random.seed(42)

# Simulate two-regime random walk: 200 calm ticks then 200 volatile ticks
_prices_raw = [0.50]
for i in range(400):
    sigma = 0.002 if i < 200 else 0.015  # calm then volatile
    ret = random.gauss(0, sigma)
    _prices_raw.append(max(0.01, _prices_raw[-1] + ret))

historical_returns = hz.prices_to_returns(_prices_raw)

# Also used by later snippets
returns = historical_returns

# Counters
_pass = 0
_fail = 0
_skip = 0


def run_snippet(name, fn, *, skip=False):
    global _pass, _fail, _skip
    if skip:
        _skip += 1
        print(f"  SKIP: {name}")
        return
    try:
        fn()
        _pass += 1
        print(f"  PASS: {name}")
    except Exception as e:
        _fail += 1
        print(f"  FAIL: {name}")
        traceback.print_exc()
        print()


# =========================================================================== #
# Snippet 1 — Quick Start: Pre-Trained Model (lines 44-61)
# hz.run() call is skipped; we test model creation + fit + markov_regime factory
# =========================================================================== #
def snippet_1_pre_trained_model():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)
    assert model.is_trained()
    assert model.n_states() == 2
    # Verify markov_regime factory accepts model kwarg (don't call hz.run)
    fn = hz.markov_regime(model=model, feed="book")
    assert callable(fn)

run_snippet("Snippet 1 - Pre-trained model (Quick Start)", snippet_1_pre_trained_model)


# =========================================================================== #
# Snippet 2 — Quick Start: Auto-Train Mode (lines 68-75)
# hz.run() call is skipped; test that markov_regime() factory with auto-train
# args returns a callable
# =========================================================================== #
def snippet_2_auto_train_mode():
    fn = hz.markov_regime(n_states=2, warmup=200, feed="book")
    assert callable(fn)

run_snippet("Snippet 2 - Auto-train mode (Quick Start)", snippet_2_auto_train_mode)


# =========================================================================== #
# Snippet 3 — Constructor (line 88)
# =========================================================================== #
def snippet_3_constructor():
    model = hz.MarkovRegimeModel(n_states=2)
    assert model.n_states() == 2

run_snippet("Snippet 3 - MarkovRegimeModel constructor", snippet_3_constructor)


# =========================================================================== #
# Snippet 4 — fit() (line 100)
# =========================================================================== #
def snippet_4_fit():
    model = hz.MarkovRegimeModel(n_states=2)
    log_likelihood = model.fit(historical_returns, max_iters=100, tol=1e-6)
    assert isinstance(log_likelihood, float)
    assert math.isfinite(log_likelihood)

run_snippet("Snippet 4 - model.fit()", snippet_4_fit)


# =========================================================================== #
# Snippet 5 — decode() (lines 116-117)
# =========================================================================== #
def snippet_5_decode():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)
    data = historical_returns
    states = model.decode(data)
    # [0, 0, 0, 1, 1, 1, 0, 0, ...]
    assert isinstance(states, list)
    assert len(states) == len(data)
    assert all(s in (0, 1) for s in states)

run_snippet("Snippet 5 - model.decode()", snippet_5_decode)


# =========================================================================== #
# Snippet 6 — filter_step() (lines 127-128)
# =========================================================================== #
def snippet_6_filter_step():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)
    observation = historical_returns[0]
    probs = model.filter_step(observation)
    # [0.85, 0.15] -> 85% probability of state 0
    assert isinstance(probs, list)
    assert len(probs) == 2
    assert abs(sum(probs) - 1.0) < 1e-6

run_snippet("Snippet 6 - model.filter_step()", snippet_6_filter_step)


# =========================================================================== #
# Snippet 7 — predict() (lines 138-139)
# =========================================================================== #
def snippet_7_predict():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)
    # Run at least one filter_step so there is a filtered state
    model.filter_step(historical_returns[0])
    next_probs = model.predict()
    # [0.78, 0.22]
    assert isinstance(next_probs, list)
    assert len(next_probs) == 2
    assert abs(sum(next_probs) - 1.0) < 1e-6

run_snippet("Snippet 7 - model.predict()", snippet_7_predict)


# =========================================================================== #
# Snippet 8 — smooth() (lines 147-148)
# =========================================================================== #
def snippet_8_smooth():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)
    data = historical_returns
    smoothed = model.smooth(data)
    # [[0.9, 0.1], [0.85, 0.15], ...]  # one row per observation
    assert isinstance(smoothed, list)
    assert len(smoothed) == len(data)
    for row in smoothed:
        assert len(row) == 2
        assert abs(sum(row) - 1.0) < 1e-6

run_snippet("Snippet 8 - model.smooth()", snippet_8_smooth)


# =========================================================================== #
# Snippet 9 — Other methods: transition_matrix, emission_params, etc.
# (lines 155-161)
# =========================================================================== #
def snippet_9_other_methods():
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)

    # transition_matrix -> list[list[float]]
    tm = model.transition_matrix()
    assert len(tm) == 2
    for row in tm:
        assert len(row) == 2
        assert abs(sum(row) - 1.0) < 1e-6

    # emission_params -> list[tuple[float, float]]
    ep = model.emission_params()
    assert len(ep) == 2
    for mean, var in ep:
        assert isinstance(mean, float)
        assert isinstance(var, float)

    # current_regime -> int (need a filter_step first)
    model.filter_step(historical_returns[0])
    cr = model.current_regime()
    assert cr in (0, 1)

    # filtered_probs -> list[float]
    fp = model.filtered_probs()
    assert len(fp) == 2
    assert abs(sum(fp) - 1.0) < 1e-6

    # reset_filter -> None
    model.reset_filter()

    # n_states -> int
    assert model.n_states() == 2

    # is_trained -> bool
    assert model.is_trained()

run_snippet("Snippet 9 - Other methods (transition_matrix, emission_params, etc.)", snippet_9_other_methods)


# =========================================================================== #
# Snippet 10 — hz.prices_to_returns (lines 170-172)
# =========================================================================== #
def snippet_10_prices_to_returns():
    prices = [100.0, 101.0, 99.0, 102.0]
    returns = hz.prices_to_returns(prices)
    # [0.00995, -0.02005, 0.02985]
    assert isinstance(returns, list)
    assert len(returns) == 3
    # Check approximate values
    assert abs(returns[0] - math.log(101.0 / 100.0)) < 1e-4
    assert abs(returns[1] - math.log(99.0 / 101.0)) < 1e-4
    assert abs(returns[2] - math.log(102.0 / 99.0)) < 1e-4

run_snippet("Snippet 10 - hz.prices_to_returns()", snippet_10_prices_to_returns)


# =========================================================================== #
# Snippet 11 — hz.markov_regime() factory (lines 184-191)
# =========================================================================== #
def snippet_11_markov_regime_factory():
    fn = hz.markov_regime(
        model=None,           # Pre-trained model (or None for auto-train)
        n_states=2,           # States for auto-train (ignored if model provided)
        warmup=100,           # Ticks before auto-training
        feed="book",          # Feed name to read prices from
        param_name="regime",  # Key in ctx.params
    )
    assert callable(fn)

run_snippet("Snippet 11 - hz.markov_regime() factory", snippet_11_markov_regime_factory)


# =========================================================================== #
# Snippet 12 — Regime-Adaptive Market Maker (lines 219-249)
# hz.run() skipped; test the strategy function logic with synthetic context
# =========================================================================== #
def snippet_12_regime_adaptive_mm():
    # Train on historical data (from doc)
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)

    # Define the adaptive_quoter exactly as in the docs
    def adaptive_quoter(ctx):
        regime = ctx.params.get("regime", 0)
        vol_prob = ctx.params.get("regime_vol_state", 0.0)

        # Base spread widens with volatility
        base_spread = 0.04
        spread = base_spread + vol_prob * 0.06  # 4-10 cents

        # Reduce size in volatile regime
        size = 10.0 if regime == 0 else 3.0

        fair = ctx.feed.price
        return hz.quotes(fair=fair, spread=spread, size=size)

    # Test with calm regime context
    class MockFeed:
        price = 0.55
        bid = 0.54
        ask = 0.56
    class MockCtx:
        params = {"regime": 0, "regime_vol_state": 0.1}
        feed = MockFeed()

    result = adaptive_quoter(MockCtx())
    assert isinstance(result, list)
    assert len(result) == 1
    q = result[0]
    assert q.size == 10.0

    # Test with volatile regime
    MockCtx.params = {"regime": 1, "regime_vol_state": 0.9}
    result2 = adaptive_quoter(MockCtx())
    assert result2[0].size == 3.0

run_snippet("Snippet 12 - Regime-Adaptive Market Maker", snippet_12_regime_adaptive_mm)


# =========================================================================== #
# Snippet 13 — Regime-Gated Trading (lines 257-270)
# hz.run() skipped; test the gate function logic
# =========================================================================== #
def snippet_13_regime_gate():
    # Define exactly as in docs
    def regime_gate(ctx):
        regime = ctx.params.get("regime", 0)
        if regime != 0:
            return None  # Skip quoting in volatile regime
        return True  # Pass through to next pipeline stage

    class MockCtx:
        params = {"regime": 0}

    # Calm regime -> pass through
    assert regime_gate(MockCtx()) is True

    # Volatile regime -> skip
    MockCtx.params = {"regime": 1}
    assert regime_gate(MockCtx()) is None

run_snippet("Snippet 13 - Regime-gated trading", snippet_13_regime_gate)


# =========================================================================== #
# Snippet 14 — Backtest with Regime Detection (lines 276-292)
# hz.backtest() and hz.run() skipped (needs engine); test model + strategy fn
# =========================================================================== #
def snippet_14_backtest_regime():
    # Train model on historical returns
    model = hz.MarkovRegimeModel(n_states=2)
    model.fit(historical_returns, max_iters=100)

    def regime_quoter(ctx):
        regime = ctx.params.get("regime", 0)
        spread = 0.06 if regime == 1 else 0.04
        return hz.quotes(fair=ctx.feed.price, spread=spread, size=5)

    # Test the quoter logic
    class MockFeed:
        price = 0.50
    class MockCtx:
        params = {"regime": 0}
        feed = MockFeed()

    result = regime_quoter(MockCtx())
    assert len(result) == 1
    assert result[0].size == 5.0
    # Calm regime: spread=0.04, so bid~0.48, ask~0.52
    assert abs(result[0].bid - 0.48) < 0.01
    assert abs(result[0].ask - 0.52) < 0.01

    # Volatile regime
    MockCtx.params = {"regime": 1}
    result2 = regime_quoter(MockCtx())
    # Volatile: spread=0.06, so bid~0.47, ask~0.53
    assert abs(result2[0].bid - 0.47) < 0.01
    assert abs(result2[0].ask - 0.53) < 0.01

run_snippet("Snippet 14 - Backtest with Regime Detection", snippet_14_backtest_regime)


# =========================================================================== #
# Snippet 15 — Offline Analysis (lines 298-321)
# =========================================================================== #
def snippet_15_offline_analysis():
    # Fit model
    model = hz.MarkovRegimeModel(n_states=3)
    model.fit(returns, max_iters=200)

    # Decode most likely state sequence
    states = model.decode(returns)
    msg = f"Regime sequence: {states[:20]}"
    assert isinstance(states, list)
    assert len(states) == len(returns)
    assert all(s in (0, 1, 2) for s in states)
    print(f"    -> {msg}")

    # Inspect learned parameters
    for i, (mean, var) in enumerate(model.emission_params()):
        print(f"    -> State {i}: mean={mean:.6f}, std={var**0.5:.6f}")

    # Transition matrix
    tm = model.transition_matrix()
    for row in tm:
        formatted = [f"{p:.3f}" for p in row]
        print(f"    -> {formatted}")

    # Verify dimensions
    assert len(tm) == 3
    for row in tm:
        assert len(row) == 3
        assert abs(sum(row) - 1.0) < 1e-6

    # Smooth for full posterior
    smoothed = model.smooth(returns)
    # smoothed[t] = [P(state_0|all_data), P(state_1|all_data), ...]
    assert len(smoothed) == len(returns)
    for row in smoothed:
        assert len(row) == 3
        assert abs(sum(row) - 1.0) < 1e-6

run_snippet("Snippet 15 - Offline Analysis", snippet_15_offline_analysis)


# =========================================================================== #
# Summary
# =========================================================================== #
print()
print("=" * 60)
print(f"  TOTAL: {_pass + _fail + _skip} snippets")
print(f"  PASS:  {_pass}")
print(f"  FAIL:  {_fail}")
print(f"  SKIP:  {_skip}")
print("=" * 60)

if _fail > 0:
    exit(1)
