"""Leakage gate: nothing in the inputs may depend on bars after the origin; the deliberate-leak fixture must be rejected."""
import numpy as np
import pytest

from chronos_x.data.synthetic import synthetic_bars
from chronos_x.data.features import build_grid, N_DAY
from chronos_x.data.windows import SymbolGrid, WindowConfig, make_example, eligible_origins


import polars as pl


def _grid(seed=3, mutate_future=None):
    """mutate_future = (s, m): every bar strictly after origin (s, m) gets its prices multiplied by 1.05 and its volume x3."""
    df, actions = synthetic_bars(start="2024-01-02", end="2024-03-15", seed=seed, p_missing=0.03)
    if mutate_future is not None:
        s, m = mutate_future
        from chronos_x.data.calendar import sessions
        sess = sessions("2024-01-02", "2024-03-15")
        cutoff = sess[s].open_utc + np.timedelta64(m, "m")      # bar-open instant of slot m
        fut = pl.col("timestamp_utc") > pl.lit(cutoff.to_pydatetime())
        df = df.with_columns([pl.when(fut).then(pl.col(c) * 1.05).otherwise(pl.col(c)).alias(c) for c in ("open", "high", "low", "close")]
                             + [pl.when(fut).then(pl.col("volume") * 3).otherwise(pl.col("volume")).alias("volume")])
    arrays, meta = build_grid(df, actions, "2024-01-02", "2024-03-15", "SYN", "stock")
    return SymbolGrid("SYN", arrays, meta.__dict__)


def test_inputs_invariant_to_future_bars():
    cfg = WindowConfig()
    g0 = _grid()
    s, m = 30, 200
    ex0 = make_example(g0, s, m, cfg)

    g1 = _grid(mutate_future=(s, m))
    ex1 = make_example(g1, s, m, cfg)
    for k in ("xA", "xB", "xC", "glob", "sigma", "fut_cal", "mA", "mB", "mC"):
        assert np.array_equal(ex0[k], ex1[k]), f"input {k} depends on the future"
    assert not np.array_equal(ex0["y"], ex1["y"])


def test_targets_identity_with_log_close():
    cfg = WindowConfig()
    g = _grid()
    for s, m in eligible_origins(g, cfg, np.array([25, 40]), stride=37)[:20]:
        ex = make_example(g, int(s), int(m), cfg)
        y, ym = ex["y"], ex["y_mask"] > 0
        for H in (5, 15, 30, 60):
            fut = np.arange(1, H + 1)
            present = ym[:H]
            if not present.any():
                continue
            last = s * N_DAY + m + fut[present][-1]
            exact = g.flat_lc[last] - g.flat_lc[s * N_DAY + m]
            assert np.isclose(y[:H][present].sum(), exact, atol=1e-6)


def test_bucket_completion_rule():
    """At origin m, Branch B's last bucket must end at a slot <= m and Branch C likewise."""
    cfg = WindowConfig()
    g = _grid()
    s = 30
    for m in (30, 33, 34, 35, 59, 60, 200):
        ex = make_example(g, s, m, cfg)
        jB = (m + 1) // 5 - 1
        assert 5 * jB + 4 <= m and 5 * (jB + 1) + 4 > m
        jC = (m + 1) // 30 - 1
        assert 30 * jC + 29 <= m and 30 * (jC + 1) + 29 > m
        # last Branch-B row equals bucket jB of session s (close feature consistent)
        assert ex["mB"][-1] == g.flat_m5[s * 78 + jB]


def test_deliberate_future_feature_fixture_is_rejected():
    """Quarantined fixture: a feature built from the next bar must be detected by the invariance check."""
    cfg = WindowConfig()
    g0 = _grid(); s, m = 30, 200
    ex0 = make_example(g0, s, m, cfg)

    def leaky_example(g):
        ex = make_example(g, s, m, cfg)
        ex["xA"] = ex["xA"].copy()
        ex["xA"][-1, 0] = g.flat_lc[s * N_DAY + m + 1] - g.flat_lc[s * N_DAY + m]   # uses t+1 close
        return ex

    g1 = _grid(mutate_future=(s, m))
    leak0, leak1 = leaky_example(g0), leaky_example(g1)
    with pytest.raises(AssertionError):
        assert np.array_equal(leak0["xA"], leak1["xA"])
    assert np.array_equal(ex0["xA"], make_example(g1, s, m, cfg)["xA"])


def test_origin_eligibility_respects_close_and_early_close():
    cfg = WindowConfig()
    df, actions = synthetic_bars(start="2024-11-20", end="2024-12-05", seed=5, p_missing=0.0)
    arrays, meta = build_grid(df, actions, "2024-11-20", "2024-12-05", "SYN", "stock")
    g = SymbolGrid("SYN", arrays, meta.__dict__)
    i = meta.session_dates.index("2024-11-29")
    o = eligible_origins(g, cfg, np.array([i, i + 1]))
    early = o[o[:, 0] == i][:, 1]; normal = o[o[:, 0] == i + 1][:, 1]
    assert early.max() + cfg.H <= 209 and normal.max() + cfg.H <= 389
    assert early.min() == cfg.m_min and normal.max() == 389 - cfg.H