from fxhnt.application.unlock_runner import UnlockShortRunner


def _series(start, n, close=100.0, qvol=1e8, funding=0.0, vol=0.03):
    out = {}
    px = close
    for i in range(n):
        # alternate up/down so name vol is well-defined and > 0
        px = px * (1.0 + (vol if i % 2 == 0 else -vol))
        out[start + i] = (px, qvol, funding)
    return out


def test_in_window_name_gets_negative_weight():
    # ALT has a big unlock cliff at day 150; window is [cliff-pre, cliff)
    panel = {
        "ALTUSDT": _series(100, 80),
        "BTCUSDT": _series(100, 80),
    }
    events = [{"sym": "ALT", "cliff_day": 150, "n": 5e7, "cat": "insiders"}]
    r = UnlockShortRunner(panel, events, pre_window=10, min_qvol=1e6, min_shock=0.0)
    w = r.current_weights("epoch:145")  # 140 <= 145 < 150 -> in window
    assert w.get("ALTUSDT", 0.0) < 0


def test_no_active_window_is_empty():
    panel = {"ALTUSDT": _series(100, 80), "BTCUSDT": _series(100, 80)}
    events = [{"sym": "ALT", "cliff_day": 150, "n": 5e7, "cat": "insiders"}]
    r = UnlockShortRunner(panel, events, pre_window=10, min_qvol=1e6, min_shock=0.0)
    assert r.current_weights("epoch:120") == {}  # well before window


def _reference_current_weights(panel, events, kwargs, as_of):
    """Recompute current_weights via the NAIVE pre-optimization logic (rebuild events_norm + _eligible
    every call + full-panel scan/sort for the latest day) so the golden sweep pins bit-identical output
    of the memoized implementation against a from-scratch baseline."""
    from fxhnt.application._date_axis import as_of_day
    from fxhnt.domain.unlock_shorts import inverse_vol_capped_weights, supply_shock

    pre = kwargs["pre_window"]
    min_qvol = kwargs["min_qvol"]
    min_shock = kwargs["min_shock"]
    vwn = kwargs.get("vol_window_name", 20)
    maxw = kwargs.get("max_name_weight", 0.25)
    events_norm = [(e["sym"] + "USDT" if not e["sym"].endswith("USDT") else e["sym"],
                    int(e["cliff_day"]), float(e["n"])) for e in events]
    elig = []
    for sym, cd, n in events_norm:
        s = panel.get(sym)
        if not s or cd not in s:
            continue
        close, qvol = s[cd][0], s[cd][1]
        if not (close > 0 and qvol and qvol >= min_qvol):
            continue
        sh = supply_shock(n, close, qvol)
        if sh is not None and sh > min_shock:
            elig.append((sym, cd))
    day_lim = as_of_day(as_of)
    days = sorted({x for s in panel.values() for x in s if x <= day_lim})
    if not days:
        return {}
    d = days[-1]
    active = [(sym, cd) for sym, cd in elig if cd - pre <= d < cd]
    import numpy as np
    vols = {}
    for sym, _cd in active:
        s = panel[sym]
        dd = [x for x in sorted(s) if x <= d][-vwn - 1:]
        if len(dd) < 5:
            continue
        rets = [s[dd[i]][0] / s[dd[i - 1]][0] - 1.0
                for i in range(1, len(dd)) if s[dd[i - 1]][0] > 0 and s[dd[i]][0] > 0]
        if len(rets) < 4:
            continue
        v = float(np.std(rets))
        if v > 0:
            vols[sym] = v
    if not vols:
        return {}
    w = inverse_vol_capped_weights(vols, maxw)
    return {sym: -wt for sym, wt in w.items()}


def test_golden_sweep_bit_identical():
    # multi-symbol panel + several events; sweep ~20 as_of dates incl. before-data, mid-stream,
    # after-last-day, and a no-active-window day. The optimized (memoized) current_weights must match
    # the from-scratch reference EXACTLY for every as_of.
    panel = {
        "ALTUSDT": _series(100, 120, vol=0.03),
        "BETAUSDT": _series(100, 120, vol=0.05),
        "GAMMAUSDT": _series(100, 120, vol=0.02),
        "BTCUSDT": _series(100, 120, vol=0.01),
    }
    events = [
        {"sym": "ALT", "cliff_day": 150, "n": 5e7, "cat": "insiders"},
        {"sym": "BETA", "cliff_day": 175, "n": 8e7, "cat": "investors"},
        {"sym": "GAMMA", "cliff_day": 200, "n": 3e7, "cat": "insiders"},
    ]
    kwargs = dict(pre_window=10, min_qvol=1e6, min_shock=0.0, vol_window_name=20,
                  max_name_weight=0.25)
    r = UnlockShortRunner(panel, events, **kwargs)
    # span: before any data (90,95), within panel (100..219 incl. windows + gaps), after last day (300)
    as_ofs = [f"epoch:{x}" for x in range(90, 220, 7)] + ["epoch:300", "epoch:80"]
    for ao in as_ofs:
        expected = _reference_current_weights(panel, events, kwargs, ao)
        assert r.current_weights(ao) == expected, ao
