"""
Leveraged trend without TrendRider's machinery: a trend filter times a
volatility target, implemented with 3x ETFs, tested on synthetic ETFs from
1990 and with parameters picked walk-forward.

    PYTHONPATH=. uv run python -m research.lab.fetch_long_history   # once
    PYTHONPATH=. uv run python -m research.lab.vol_trend

Exposure to the index (SPY or QQQ, total return) decided on closes through
t-1:
    E_t = 1{index > its n-day average} x min(cap, target / sigma_t)
sigma_t is the index's EWMA vol (20-day half-life). E/3 is held in the 3x
ETF (UPRO / TQQQ, synthetic: research/lab/letf_history.py), the rest in the
2-year Treasury series. Holdings drift; the book trades back to target when
the signal flips or E is more than 25% off; 10 bp per unit traded. Fills:
"lag" trades at the close after the decision (pessimistic), "close" at the
decision close (optimistic).

Moreira & Muir (2017) and Harvey et al. (2018) motivate the vol target;
Gayed & Bilello (2016) the moving-average switch. The grid is small and
round; the walk-forward chain picks, each January, the configuration with
the best Sharpe over the previous ten years.
"""

from __future__ import annotations

import itertools
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.letf_history import BLOCKS, SRC, build

warnings.filterwarnings("ignore")
OUT = os.path.join("data", "research", "results", "vol_trend.pkl")
COST, BAND = 0.001, 0.25
GRID = list(itertools.product(("SPY", "QQQ"), (100, 150, 200), (0.15, 0.20, 0.25, 0.30, 0.40), (1.5, 3.0)))
LETF = {"SPY": "UPRO", "QQQ": "TQQQ"}


def backtest(syn: pd.DataFrame, index: str, n: int, target: float, cap: float,
             lag: bool = True) -> tuple[pd.Series, pd.Series]:
    px = syn[index]
    r = px.pct_change().fillna(0.0)
    on = (px > px.rolling(n).mean()).astype(float)
    sigma = np.sqrt((r ** 2).ewm(halflife=20, min_periods=40).mean() * 252)
    E = (on * (target / sigma).clip(upper=cap)).fillna(0.0)
    E.iloc[:n] = 0.0
    tgt = (E / 3.0).shift(2 if lag else 1).fillna(0.0).to_numpy()   # row t = weight held over t
    rl = syn[LETF[index]].pct_change().fillna(0.0).to_numpy()
    rs = syn["SHY"].pct_change().fillna(0.0).to_numpy()
    w, out, held = 0.0, np.empty(len(tgt)), np.empty(len(tgt))
    for t in range(len(tgt)):
        if t > 0:
            g = w * (1 + rl[t]) + (1 - w) * (1 + rs[t])
            out[t] = g - 1
            w = w * (1 + rl[t]) / g
        else:
            out[t] = 0.0
        want = tgt[t + 1] if t + 1 < len(tgt) else w
        if (want == 0) != (w == 0) or abs(want - w) > BAND * max(w, 1e-9):
            out[t] -= abs(want - w) * COST * 2              # sell one leg, buy the other
            w = want
        held[t] = w * 3
    return pd.Series(out, index=syn.index), pd.Series(held, index=syn.index)


def sharpe(r: pd.Series) -> float:
    return float(r.mean() / r.std() * np.sqrt(252)) if r.std() > 0 else 0.0


def block_line(r: pd.Series) -> str:
    cells = []
    for a, b in BLOCKS:
        x = r.loc[a:b]
        e = (1 + x).cumprod()
        cells.append(f"{(e.iloc[-1] ** (252 / len(x)) - 1)*100:6.1f}%/{sharpe(x):4.2f}/{((e / e.cummax()) - 1).min()*100:4.0f}%")
    return " ".join(c.rjust(17) for c in cells)


def walk_forward(R: pd.DataFrame, lookback_years: int = 10, start: int = 2000) -> tuple[pd.Series, dict]:
    parts, picks = [], {}
    for y in range(start, R.index[-1].year + 1):
        hist = R.loc[f"{y - lookback_years}-01-01":f"{y - 1}-12-31"]
        best = hist.apply(sharpe).idxmax()
        picks[y] = best
        parts.append(R.loc[f"{y}-01-01":f"{y}-12-31", best])
    return pd.concat(parts), picks


def main() -> None:
    with open(SRC, "rb") as fh:
        syn, _ = build(pickle.load(fh))
    syn = syn.loc["1989-01-01":]
    R, H, params = {}, {}, {}
    for idx, n, tv, cap in GRID:
        key = f"{idx} ma{n} vt{int(tv*100)} cap{cap:g}"
        params[key] = (idx, n, tv, cap)
        R[key], H[key] = backtest(syn, idx, n, tv, cap)
    R = pd.DataFrame(R).loc["1990-01-01":]
    head = " ".join(f"{a[:4]}-{b[2:4]:>2s}".rjust(17) for a, b in BLOCKS)
    full = R.apply(sharpe).sort_values(ascending=False)
    print(f"CAGR/Sharpe/MaxDD by block, lag fills, synthetic 3x ETFs\n{'':28s}{head}")
    for k in list(full.index[:8]) + list(full.index[len(full) // 2:len(full) // 2 + 2]) + list(full.index[-2:]):
        print(f"{k:28s}{block_line(R[k])}  avg E {H[k].loc['1990':].mean():.2f}")
    print(f"\nfull-period Sharpe over the {len(full)}-config grid: median {full.median():.2f}, "
          f"range {full.min():.2f}-{full.max():.2f}")
    wf, picks = walk_forward(R)
    print(f"\nwalk-forward (best trailing-10y Sharpe each January), 2000-2026:\n{'':28s}{head}")
    print(f"{'walk-forward chain':28s}{block_line(pd.concat([R.iloc[:, 0].loc[:'1999'] * np.nan, wf]).fillna(0.0))}")
    print("picks:", {y: p for y, p in picks.items() if y in (2000, 2005, 2010, 2015, 2020, 2025, 2026)})
    close_mode = {k: backtest(syn, *params[k], lag=False)[0].loc["1990-01-01":] for k in full.index[:3]}
    print("\nsame top configs with close fills (optimistic):")
    for k, r in close_mode.items():
        print(f"{k:28s}{block_line(r)}")
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "wb") as fh:
        pickle.dump({"returns": R, "exposure": pd.DataFrame(H), "walk_forward": wf, "picks": picks}, fh)


if __name__ == "__main__":
    main()