"""
Tick-volume normalization test. Pre-registered in VOLNORM_PLAN.md.

    python research/volnorm.py [H4|H1]
"""

from __future__ import annotations

import os
import sys

import numpy as np
import pandas as pd
from scipy.stats import spearmanr

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import afml_bars as ab   # noqa: E402  (m1(), agg(), HCC path)

SLOT_WINDOW = 60          # sessions of the same hour-of-day slot
Z_ENTRY, Z_WIN, MAX_BARS = -1.5, 20, 10
BOOT, BLOCK = 5000, 10
rng = np.random.default_rng(7)


def bars(symbol: str, secs: int) -> pd.DataFrame:
    a = ab.m1(symbol)
    b = ab.agg(a, a["t"] // secs, 0.0, 0.0)
    df = pd.DataFrame({k: b[k] for k in ("ts", "o", "h", "l", "c", "v")})
    df = df[df.v > 0].reset_index(drop=True)
    df["hour"] = pd.DatetimeIndex(df.ts).hour
    df["year"] = pd.DatetimeIndex(df.ts).year
    return df


def old_rvol(v: pd.Series) -> pd.Series:
    """What the EA computes today: bar / mean of the previous 20 bars (shift excludes the bar itself
    from the mean only for causality of the comparison; the EA includes it - immaterial to the R^2)."""
    return np.log(v / v.rolling(20).mean())


def new_rvol(df: pd.DataFrame) -> pd.Series:
    prof = df.groupby("hour")["v"].transform(
        lambda s: s.shift(1).rolling(SLOT_WINDOW, min_periods=SLOT_WINDOW // 2).median())
    return np.log(df.v / prof)


def r2_on_hour(x: pd.Series, hour: pd.Series) -> float:
    ok = x.notna() & np.isfinite(x)
    x, h = x[ok], hour[ok]
    tot = ((x - x.mean()) ** 2).sum()
    within = ((x - x.groupby(h).transform("mean")) ** 2).sum()
    return 1.0 - within / tot


def year_spread(x: pd.Series, year: pd.Series) -> float:
    return float(x.groupby(year).mean().std())


def atr14(df: pd.DataFrame) -> pd.Series:
    pc = df.c.shift(1)
    tr = pd.concat([df.h - df.l, (df.h - pc).abs(), (df.l - pc).abs()], axis=1).max(axis=1)
    return tr.rolling(14).mean()


def dip_events(df: pd.DataFrame) -> pd.DataFrame:
    """z20 <= -1.5 on the closed bar; fill next open; exit at first close >= SMA20 else 10 bars."""
    sma = df.c.rolling(Z_WIN).mean()
    sd = df.c.rolling(Z_WIN).std()
    z = (df.c - sma) / sd
    rows, i, n = [], Z_WIN, len(df)
    c, o, s = df.c.values, df.o.values, sma.values
    while i < n - MAX_BARS - 2:
        if z.iloc[i] <= Z_ENTRY:
            entry = o[i + 1]
            j = i + 1
            while j < min(i + 1 + MAX_BARS, n - 1) and c[j] < s[j]:
                j += 1
            rows.append((i, (c[j] / entry - 1.0) * 1e4))
            i = j + 1
        else:
            i += 1
    return pd.DataFrame(rows, columns=["i", "bp"])


def terciles(feat: np.ndarray, bp: np.ndarray):
    lo, hi = np.nanpercentile(feat, [33.3, 66.7])
    return bp[feat <= lo], bp[feat >= hi]


def block_boot_diff(top: np.ndarray, bot: np.ndarray) -> tuple[float, float, float]:
    d = top.mean() - bot.mean()
    ds = np.empty(BOOT)
    for k in range(BOOT):
        def rs(x):
            nb = max(1, len(x) // BLOCK)
            st = rng.integers(0, max(1, len(x) - BLOCK), nb)
            return np.concatenate([x[s:s + BLOCK] for s in st]).mean()
        ds[k] = rs(top) - rs(bot)
    return d, *np.percentile(ds, [2.5, 97.5])


def main() -> None:
    label = sys.argv[1] if len(sys.argv) > 1 else "H4"
    secs = {"H4": 4 * 3600, "H1": 3600}[label]
    print(f"== bars: {label}, slot window {SLOT_WINDOW} sessions ==\n")
    print(f"{'idx':7s} {'R2 old':>7s} {'R2 new':>7s} {'yr sd old':>10s} {'yr sd new':>10s}"
          f" {'rho old':>8s} {'rho new':>8s}")
    pooled = {"old": ([], []), "new": ([], [])}
    per_idx = []
    for sym in ab.IDX:
        df = bars(sym, secs)
        df["old"], df["new"] = old_rvol(df.v), new_rvol(df)
        atr = atr14(df)
        act = (df.c.shift(-1) / df.c - 1.0).abs() * df.c / atr
        m = df.old.notna() & df.new.notna() & act.notna()
        rho_o = spearmanr(df.old[m], act[m]).statistic
        rho_n = spearmanr(df.new[m], act[m]).statistic
        print(f"{sym:7s} {r2_on_hour(df.old, df.hour):7.3f} {r2_on_hour(df.new, df.hour):7.3f}"
              f" {year_spread(df.old, df.year):10.3f} {year_spread(df.new, df.year):10.3f}"
              f" {rho_o:8.3f} {rho_n:8.3f}")
        ev = dip_events(df)
        ev = ev[df.old.iloc[ev.i].notna().values & df.new.iloc[ev.i].notna().values]
        res = {}
        for kind in ("old", "new"):
            f = df[kind].iloc[ev.i].values
            top, bot = terciles(f, ev.bp.values)
            res[kind] = block_boot_diff(top, bot)
            pooled[kind][0].append(top)
            pooled[kind][1].append(bot)
        per_idx.append((sym, len(ev), ev.bp.mean(), res))
    print("\nDip-z events: mean bp, then top-minus-bottom tercile of the volume feature "
          "(diff [95% block-bootstrap CI])")
    for sym, n, mbp, res in per_idx:
        print(f"{sym:7s} n={n:4d} mean {mbp:6.1f} bp | "
              + " | ".join(f"{k}: {d:6.1f} [{lo:6.1f},{hi:6.1f}]" for k, (d, lo, hi) in res.items()))
    for kind in ("old", "new"):
        top = np.concatenate(pooled[kind][0])
        bot = np.concatenate(pooled[kind][1])
        d, lo, hi = block_boot_diff(top, bot)
        print(f"POOLED {kind}: top {top.mean():6.1f} bp (n {len(top)}), bottom {bot.mean():6.1f} bp"
              f" (n {len(bot)}), diff {d:6.1f} [{lo:6.1f}, {hi:6.1f}]")


if __name__ == "__main__":
    main()