"""Test every Python code snippet from docs/bars-labeling.mdx.

Each snippet is wrapped in a try/except that prints PASS/FAIL.
Synthetic data is used where the doc uses placeholder [...].
"""

import math
import random
from collections import Counter

import horizon as hz

# ---------------------------------------------------------------------------
# Synthetic data generation
# ---------------------------------------------------------------------------
random.seed(42)
N_TICKS = 2000

# Generate a realistic random walk for tick data
_base_price = 100.0
_prices_raw = [_base_price]
for _ in range(N_TICKS - 1):
    _prices_raw.append(_prices_raw[-1] * (1 + random.gauss(0, 0.002)))

timestamps = [float(i) for i in range(N_TICKS)]
prices = _prices_raw
volumes = [random.uniform(0.5, 10.0) for _ in range(N_TICKS)]

# ---------------------------------------------------------------------------
# Test execution
# ---------------------------------------------------------------------------
passed = 0
failed = 0
results = []


def run_test(name, fn):
    global passed, failed
    try:
        fn()
        print(f"  PASS: {name}")
        results.append(("PASS", name))
        passed += 1
    except Exception as e:
        print(f"  FAIL: {name} -- {e}")
        results.append(("FAIL", name, str(e)))
        failed += 1


# ===========================================================================
# Snippet 1: Tick Bars (Standard Bars tab)
# ===========================================================================
def test_tick_bars():
    bars = hz.tick_bars(timestamps, prices, volumes, threshold=100)
    assert isinstance(bars, list)
    assert len(bars) > 0
    for b in bars:
        assert hasattr(b, "open")
        assert hasattr(b, "close")

run_test("Snippet 1: hz.tick_bars()", test_tick_bars)


# ===========================================================================
# Snippet 2: Volume Bars (Standard Bars tab)
# ===========================================================================
def test_volume_bars():
    bars = hz.volume_bars(timestamps, prices, volumes, threshold=1000.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 2: hz.volume_bars()", test_volume_bars)


# ===========================================================================
# Snippet 3: Dollar Bars (Standard Bars tab)
# ===========================================================================
def test_dollar_bars():
    bars = hz.dollar_bars(timestamps, prices, volumes, threshold=50000.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 3: hz.dollar_bars()", test_dollar_bars)


# ===========================================================================
# Snippet 4: Tick Imbalance Bars (Imbalance Bars tab)
# ===========================================================================
def test_tick_imbalance_bars():
    bars = hz.tick_imbalance_bars(timestamps, prices, volumes, initial_estimate=50.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 4: hz.tick_imbalance_bars()", test_tick_imbalance_bars)


# ===========================================================================
# Snippet 5: Volume Imbalance Bars (Imbalance Bars tab)
# ===========================================================================
def test_volume_imbalance_bars():
    bars = hz.volume_imbalance_bars(timestamps, prices, volumes, initial_estimate=500.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 5: hz.volume_imbalance_bars()", test_volume_imbalance_bars)


# ===========================================================================
# Snippet 6: Tick Run Bars (Run Bars tab)
# ===========================================================================
def test_tick_run_bars():
    bars = hz.tick_run_bars(timestamps, prices, volumes, initial_estimate=50.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 6: hz.tick_run_bars()", test_tick_run_bars)


# ===========================================================================
# Snippet 7: Volume Run Bars (Run Bars tab)
# ===========================================================================
def test_volume_run_bars():
    bars = hz.volume_run_bars(timestamps, prices, volumes, initial_estimate=500.0)
    assert isinstance(bars, list)
    assert len(bars) > 0

run_test("Snippet 7: hz.volume_run_bars()", test_volume_run_bars)


# ===========================================================================
# Snippet 8: Comparing Bar Types (full example)
# ===========================================================================
def test_comparing_bar_types():
    # Standard: fixed 100-tick bars
    tick = hz.tick_bars(timestamps, prices, volumes, threshold=100)

    # Information-driven: bars adapt to market activity
    tib = hz.tick_imbalance_bars(timestamps, prices, volumes, initial_estimate=50.0)
    trb = hz.tick_run_bars(timestamps, prices, volumes, initial_estimate=50.0)

    print(f"    Tick bars: {len(tick)}")
    print(f"    Tick imbalance bars: {len(tib)}")
    print(f"    Tick run bars: {len(trb)}")

    assert len(tick) > 0
    assert len(tib) > 0
    assert len(trb) > 0

    # Inspect a bar
    bar = tick[0]
    print(f"    O={bar.open}, H={bar.high}, L={bar.low}, C={bar.close}")
    print(f"    Volume={bar.volume}, VWAP={bar.vwap}, Ticks={bar.n_ticks}")

    assert bar.high >= bar.low
    assert bar.n_ticks > 0
    assert bar.volume > 0
    assert bar.vwap > 0

run_test("Snippet 8: Comparing Bar Types", test_comparing_bar_types)


# ===========================================================================
# Snippet 9: get_daily_vol (Step 1)
# ===========================================================================
def test_get_daily_vol():
    # Exponentially weighted daily volatility from log returns
    vol = hz.get_daily_vol(prices, span=20)
    assert isinstance(vol, list)
    assert len(vol) == len(prices)
    assert vol[0] == 0.0  # first element is always 0.0

run_test("Snippet 9: hz.get_daily_vol()", test_get_daily_vol)


# ===========================================================================
# Snippet 10: cusum_filter (Step 2)
# ===========================================================================
def test_cusum_filter():
    events = hz.cusum_filter(prices, threshold=0.02)
    assert isinstance(events, list)
    # Should detect some events in our random walk
    # (threshold 0.02 in price units is very small for prices ~100,
    #  so we should get many events)
    assert len(events) > 0
    # All event indices should be valid
    for idx in events:
        assert 0 <= idx < len(prices)

run_test("Snippet 10: hz.cusum_filter()", test_cusum_filter)


# ===========================================================================
# Snippet 11: triple_barrier_labels (Step 3)
# ===========================================================================
def test_triple_barrier_labels():
    events = hz.cusum_filter(prices, threshold=0.5)
    # If threshold is too large and no events, try smaller
    if len(events) == 0:
        events = hz.cusum_filter(prices, threshold=0.02)
    assert len(events) > 0, "Need at least one event for triple barrier labeling"

    labels = hz.triple_barrier_labels(
        prices=prices,
        timestamps=timestamps,
        events=events,
        pt_sl=[1.0, 1.0],
        min_ret=0.005,
        max_holding=100,
        vol_span=20,
    )
    assert isinstance(labels, list)
    assert len(labels) == len(events)
    for lbl in labels:
        assert hasattr(lbl, "event_idx")
        assert hasattr(lbl, "label")
        assert hasattr(lbl, "ret")
        assert hasattr(lbl, "barrier")
        assert hasattr(lbl, "touch_idx")
        assert lbl.label in (-1, 0, 1)
        assert lbl.barrier in ("pt", "sl", "vb")

run_test("Snippet 11: hz.triple_barrier_labels()", test_triple_barrier_labels)


# ===========================================================================
# Snippet 12: meta_labels (Step 4)
# ===========================================================================
def test_meta_labels():
    meta = hz.meta_labels(
        prices=prices,
        timestamps=timestamps,
        primary_signals=[(10, 1), (50, -1), (120, 1)],
        pt_sl=[2.0, 1.0],
        max_holding=50,
        vol_span=20,
    )
    assert isinstance(meta, list)
    assert len(meta) == 3
    for lbl in meta:
        assert hasattr(lbl, "event_idx")
        assert hasattr(lbl, "label")
        assert lbl.label in (0, 1)
        assert hasattr(lbl, "ret")
        assert hasattr(lbl, "barrier")
        assert hasattr(lbl, "touch_idx")

run_test("Snippet 12: hz.meta_labels()", test_meta_labels)


# ===========================================================================
# Snippet 13: drop_labels (Step 5)
# ===========================================================================
def test_drop_labels():
    # First generate some labels
    events = hz.cusum_filter(prices, threshold=0.02)
    if len(events) < 5:
        events = list(range(10, min(len(prices) - 100, 200), 20))
    labels = hz.triple_barrier_labels(
        prices=prices,
        timestamps=timestamps,
        events=events,
        pt_sl=[1.0, 1.0],
        min_ret=0.005,
        max_holding=100,
        vol_span=20,
    )
    assert len(labels) > 0

    cleaned = hz.drop_labels(labels, min_pct=0.05)
    assert isinstance(cleaned, list)
    # cleaned should be a subset of labels (or equal)
    assert len(cleaned) <= len(labels)

run_test("Snippet 13: hz.drop_labels()", test_drop_labels)


# ===========================================================================
# Snippet 14: Full Pipeline Example
# ===========================================================================
def test_full_pipeline():
    # 1. Build information-driven bars
    bars = hz.volume_bars(timestamps, prices, volumes, threshold=10000.0)
    # If threshold too high, use smaller
    if len(bars) < 10:
        bars = hz.volume_bars(timestamps, prices, volumes, threshold=500.0)
    assert len(bars) > 2, f"Need at least 3 bars, got {len(bars)}"

    bar_prices = [b.close for b in bars]
    bar_timestamps = [b.timestamp for b in bars]

    # 2. Compute daily volatility
    vol = hz.get_daily_vol(bar_prices, span=20)
    assert len(vol) == len(bar_prices)

    # 3. CUSUM filter for structural breaks
    events = hz.cusum_filter(bar_prices, threshold=0.02)
    print(f"    CUSUM detected {len(events)} events")

    # If no events detected with 0.02, try smaller threshold
    if len(events) == 0:
        events = hz.cusum_filter(bar_prices, threshold=0.001)
        print(f"    CUSUM (retry) detected {len(events)} events")

    if len(events) == 0:
        # Use manual events if CUSUM finds nothing
        events = list(range(1, min(len(bar_prices) - 2, 10)))
        print(f"    Using manual events: {len(events)} events")

    # 4. Triple barrier labeling
    labels = hz.triple_barrier_labels(
        prices=bar_prices,
        timestamps=bar_timestamps,
        events=events,
        pt_sl=[1.0, 1.0],
        min_ret=0.005,
        max_holding=100,
        vol_span=20,
    )
    assert len(labels) > 0

    # 5. Inspect label distribution
    dist = Counter(l.label for l in labels)
    print(f"    Label distribution: {dict(dist)}")

    # 6. Drop rare labels (< 5%)
    cleaned = hz.drop_labels(labels, min_pct=0.05)
    print(f"    After dropping rare labels: {len(cleaned)} / {len(labels)}")

    # 7. Use labels for ML training
    features = []
    targets = []
    for label in cleaned:
        idx = label.event_idx
        features.append(bar_prices[idx])
        targets.append(label.label)

    assert len(features) == len(targets)
    assert len(features) == len(cleaned)

run_test("Snippet 14: Full Pipeline Example", test_full_pipeline)


# ===========================================================================
# Bar field validation (from "Bar Type" table in docs)
# ===========================================================================
def test_bar_fields():
    bars = hz.tick_bars(timestamps, prices, volumes, threshold=50)
    assert len(bars) > 0
    b = bars[0]
    # Validate all documented fields exist and have correct types
    assert isinstance(b.timestamp, float), "timestamp should be float"
    assert isinstance(b.open, float), "open should be float"
    assert isinstance(b.high, float), "high should be float"
    assert isinstance(b.low, float), "low should be float"
    assert isinstance(b.close, float), "close should be float"
    assert isinstance(b.volume, float), "volume should be float"
    assert isinstance(b.vwap, float), "vwap should be float"
    assert isinstance(b.n_ticks, int), "n_ticks should be int"

run_test("Bar field types match docs", test_bar_fields)


# ===========================================================================
# BarrierLabel field validation (from "BarrierLabel Type" table in docs)
# ===========================================================================
def test_barrier_label_fields():
    events = hz.cusum_filter(prices, threshold=0.02)
    if len(events) == 0:
        events = [10, 50, 100]
    labels = hz.triple_barrier_labels(
        prices=prices,
        timestamps=timestamps,
        events=events[:5],
        pt_sl=[1.0, 1.0],
        min_ret=0.005,
        max_holding=100,
        vol_span=20,
    )
    assert len(labels) > 0
    lbl = labels[0]
    # Validate all documented fields exist and have correct types
    assert isinstance(lbl.event_idx, int), "event_idx should be int"
    assert isinstance(lbl.label, int), "label should be int"
    assert isinstance(lbl.ret, float), "ret should be float"
    assert isinstance(lbl.barrier, str), "barrier should be str"
    assert isinstance(lbl.touch_idx, int), "touch_idx should be int"

run_test("BarrierLabel field types match docs", test_barrier_label_fields)


# ===========================================================================
# Summary
# ===========================================================================
print("\n" + "=" * 60)
print(f"SUMMARY: {passed} passed, {failed} failed, {passed + failed} total")
print("=" * 60)
for r in results:
    status = r[0]
    name = r[1]
    extra = f" -- {r[2]}" if len(r) > 2 else ""
    print(f"  [{status}] {name}{extra}")
print("=" * 60)

if failed > 0:
    exit(1)
