#!/usr/bin/env python3
"""
Backtest dollar-neutral long-short strategies from Ken French momentum and variance deciles.

These are **portfolio-return** strategies (not stock-level intersection corners). French
publishes decile portfolio returns; we combine legs (e.g. Hi PRIOR − Hi 10 variance).

Presets:
  ``all`` — classic spreads plus vol-target variants
  ``vol_target_mom_hi`` — vol-scaled Hi PRIOR − Hi 10 (default 10% ann vol)
  ``optimized`` — baseline + 24m lookback + ST-reversal blends
  ``low_dd`` — minimum drawdown variants (8% vol, blends, DD cut overlay)
  ``invested`` — always on after vol warm-up (no DD cash cut); default for live-style use

Example::

    cd /Users/robzingale/trading_bot
    PYTHONUNBUFFERED=1 .venv/bin/python \\
        RenTech/strategy_stack/run_french_mom_vol_decile_backtest.py \\
        --preset vol_target_mom_hi --target-vol 0.10 \\
        --start 1963-07-01 --capital 100000 \\
        --out-prefix RenTech/data/logs/french_mom_vol_hi_hi_vol10
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

_REPO = Path(__file__).resolve().parents[2]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from RenTech.strategy_stack.french_decile_loader import (
    FRENCH_DIR,
    load_momentum_deciles,
    load_short_term_reversal_deciles,
    load_variance_deciles,
    pct_to_decimal,
    drawdown_exposure_overlay,
    vol_scale_monthly,
)

LOGS = _REPO / "RenTech" / "data" / "logs"

STRATEGIES = {
    "mom_ls": ("Classic momentum", "Hi PRIOR", "Lo PRIOR"),
    "mom_long_hi_vol_short": ("Long hi mom, short hi vol (approx)", "Hi PRIOR", "Hi 10"),
    "mom_long_lo_vol_short": ("Long hi mom, short lo vol", "Hi PRIOR", "Lo 10"),
    "lowvol_ls": ("Low vol long, high vol short", "Lo 10", "Hi 10"),
}


def _align_frames(
    frames: list[pd.DataFrame],
    start: str | None,
    end: str | None,
) -> list[pd.DataFrame]:
    idx = frames[0].index
    for f in frames[1:]:
        idx = idx.intersection(f.index)
    out = [f.loc[idx].copy() for f in frames]
    if start:
        ts = pd.Timestamp(start)
        out = [f.loc[f.index >= ts] for f in out]
    if end:
        te = pd.Timestamp(end)
        out = [f.loc[f.index <= te] for f in out]
    return out


def _add_vol_target(
    strategy_rets: dict[str, pd.Series],
    scales_map: dict[str, pd.Series],
    raw_map: dict[str, pd.Series],
    meta: dict[str, dict],
    key: str,
    raw: pd.Series,
    *,
    description: str,
    target_vol: float,
    lookback: int,
) -> None:
    scaled, scales = vol_scale_monthly(raw, target_ann=target_vol, lookback=lookback)
    strategy_rets[key] = scaled
    scales_map[key] = scales
    raw_map[key] = raw
    meta[key] = {
        "description": description,
        "target_vol_ann": target_vol,
        "vol_lookback_months": lookback,
    }


def _leg_spread(
    mom_dec: dict[str, pd.Series],
    var_dec: dict[str, pd.Series],
    long_col: str,
    short_col: str,
) -> pd.Series:
    if long_col in mom_dec and short_col in mom_dec:
        return mom_dec[long_col] - mom_dec[short_col]
    if long_col in mom_dec and short_col in var_dec:
        return mom_dec[long_col] - var_dec[short_col]
    if long_col in var_dec and short_col in var_dec:
        return var_dec[long_col] - var_dec[short_col]
    raise KeyError(f"Columns missing: long={long_col} short={short_col}")


def _stats(monthly_ret: pd.Series, capital: float) -> dict:
    r = monthly_ret.dropna().astype(float)
    if len(r) < 12:
        return {"n_months": len(r), "sharpe_monthly": float("nan")}
    eq = capital * (1.0 + r).cumprod()
    peak = eq.cummax()
    dd = eq / peak - 1.0
    sd = r.std(ddof=1)
    sharpe = float(r.mean() / sd * np.sqrt(12.0)) if sd > 1e-12 else float("nan")
    years = len(r) / 12.0
    total_ret = float(eq.iloc[-1] / capital - 1.0)
    cagr = float((eq.iloc[-1] / capital) ** (1.0 / years) - 1.0) if years > 0 else float("nan")
    return {
        "n_months": int(len(r)),
        "start": str(r.index.min().date()),
        "end": str(r.index.max().date()),
        "capital_usd": capital,
        "ending_equity_usd": round(float(eq.iloc[-1]), 2),
        "total_return_pct": round(100.0 * total_ret, 4),
        "cagr_pct": round(100.0 * cagr, 4),
        "sharpe_monthly": round(sharpe, 4),
        "max_drawdown_pct": round(100.0 * float(dd.min()), 4),
        "vol_ann_pct": round(100.0 * float(r.std(ddof=1) * np.sqrt(12.0)), 4),
    }


def _equity_frame(
    series: pd.Series,
    scales: pd.Series | None,
    raw: pd.Series | None,
    capital: float,
    strategy: str,
) -> pd.DataFrame:
    r = series.astype(float)
    eq = capital * (1.0 + r).cumprod()
    df = pd.DataFrame({
        "date": series.index.strftime("%Y-%m-%d"),
        "strategy": strategy,
        "monthly_return": r.values,
        "monthly_pnl_usd": (r * capital).values,
        "equity_usd": eq.values,
    })
    if raw is not None:
        df["raw_spread_return"] = raw.reindex(series.index).values
    if scales is not None:
        df["vol_scale"] = scales.reindex(series.index).values
    return df


def _yearly_chained(monthly_ret: pd.Series, capital: float) -> pd.DataFrame:
    """Calendar-year compounded returns with chained equity."""
    s = monthly_ret.dropna().astype(float).copy()
    s.index = pd.to_datetime(s.index)
    cap = capital
    rows = []
    for year, g in s.groupby(s.index.year):
        r = float((1.0 + g).prod() - 1.0)
        end = cap * (1.0 + r)
        rows.append({
            "year": int(year),
            "return_pct": round(100.0 * r, 2),
            "start_equity_usd": round(cap, 2),
            "end_equity_usd": round(end, 2),
            "pnl_usd": round(end - cap, 2),
        })
        cap = end
    return pd.DataFrame(rows)


def _invested_stats(monthly_ret: pd.Series, warm: int = 24) -> dict:
    s = monthly_ret.astype(float)
    active = s.abs() > 1e-8
    return {
        "pct_months_invested_full_sample": round(100.0 * active.mean(), 1),
        "pct_months_invested_after_warmup": round(100.0 * active.iloc[warm:].mean(), 1),
    }


def run_backtest(
    *,
    start: str | None,
    end: str | None,
    capital: float,
    weighting: str,
    preset: str,
    target_vol: float,
    vol_lookback: int,
    out_prefix: Path,
    verbose: bool,
) -> dict:
    mom = load_momentum_deciles(weighting=weighting)  # type: ignore[arg-type]
    var = load_variance_deciles(weighting=weighting)  # type: ignore[arg-type]
    st = load_short_term_reversal_deciles(weighting=weighting)  # type: ignore[arg-type]
    mom, var, st = _align_frames([mom, var, st], start, end)

    mom_dec = {c: pct_to_decimal(mom[c]) for c in mom.columns}
    var_dec = {c: pct_to_decimal(var[c]) for c in var.columns}
    st_dec = {c: pct_to_decimal(st[c]) for c in st.columns}

    rho_lo_mom_hi_vol = float(
        pd.DataFrame({"lo_mom": mom_dec["Lo PRIOR"], "hi_var": var_dec["Hi 10"]}).corr().iloc[0, 1]
    )

    strategy_rets: dict[str, pd.Series] = {}
    strategy_meta_extra: dict[str, dict] = {}
    scales_map: dict[str, pd.Series] = {}
    raw_map: dict[str, pd.Series] = {}

    raw_hi_hi = _leg_spread(mom_dec, var_dec, "Hi PRIOR", "Hi 10")
    st_rev = st_dec["Lo PRIOR"] - st_dec["Hi PRIOR"]

    if preset in ("all", "vol_target_mom_hi", "optimized"):
        tvs = [target_vol] if preset in ("vol_target_mom_hi", "optimized") else [0.08, 0.10, 0.12]
        for tv in tvs:
            _add_vol_target(
                strategy_rets, scales_map, raw_map, strategy_meta_extra,
                f"vol_target_mom_hi_{int(tv * 100)}",
                raw_hi_hi,
                description=f"Vol-target Hi PRIOR − Hi 10 @ {tv:.0%} ann vol",
                target_vol=tv,
                lookback=vol_lookback,
            )

    if preset == "optimized":
        _add_vol_target(
            strategy_rets, scales_map, raw_map, strategy_meta_extra,
            "vol_target_mom_hi_10_lb24",
            raw_hi_hi,
            description="Hi PRIOR − Hi 10 @ 10% vol, 24m lookback",
            target_vol=target_vol,
            lookback=24,
        )
        vt_mom = strategy_rets[f"vol_target_mom_hi_{int(target_vol * 100)}"]
        vt_st, _ = vol_scale_monthly(st_rev, target_ann=target_vol, lookback=vol_lookback)
        for w_pct, w in ((75, 0.75), (50, 0.50)):
            key = f"vol_blend_{w_pct}_mom_{int(target_vol*100)}_strev"
            blend = w * vt_mom + (1.0 - w) * vt_st
            strategy_rets[key] = blend
            raw_map[key] = w * raw_hi_hi + (1.0 - w) * st_rev
            strategy_meta_extra[key] = {
                "description": f"{w:.0%} vol-target mom_hi_hi10 + {1-w:.0%} vol-target ST rev",
                "target_vol_ann": target_vol,
                "vol_lookback_months": vol_lookback,
            }

    if preset == "invested":
        _add_vol_target(
            strategy_rets, scales_map, raw_map, strategy_meta_extra,
            "invested_mom_10_lb24",
            raw_hi_hi,
            description="Hi PRIOR − Hi 10 @ 10% vol, 24m lookback (usually invested)",
            target_vol=0.10,
            lookback=24,
        )
        vt_mom_8_24, _ = vol_scale_monthly(raw_hi_hi, target_ann=0.08, lookback=24)
        vt_st_8_24, _ = vol_scale_monthly(st_rev, target_ann=0.08, lookback=24)
        vt_mom_10_36, _ = vol_scale_monthly(raw_hi_hi, target_ann=0.10, lookback=36)
        vt_st_10_36, _ = vol_scale_monthly(st_rev, target_ann=0.10, lookback=36)
        for key, ser, desc in (
            ("invested_mom_8_24", vt_mom_8_24, "Mom spread @ 8% vol / 24m"),
            (
                "invested_blend50_8_24",
                0.5 * vt_mom_8_24 + 0.5 * vt_st_8_24,
                "50% mom + 50% ST rev @ 8% vol / 24m (lower DD)",
            ),
            (
                "invested_blend75_10_strev",
                0.75 * vt_mom_10_36 + 0.25 * vt_st_10_36,
                "75% mom + 25% ST rev @ 10% vol / 36m",
            ),
        ):
            strategy_rets[key] = ser
            strategy_meta_extra[key] = {"description": desc}

    if preset == "low_dd":
        # Reference: original 10% / 36m
        _add_vol_target(
            strategy_rets, scales_map, raw_map, strategy_meta_extra,
            "ref_vol_target_mom_hi_10",
            raw_hi_hi,
            description="Reference: Hi PRIOR − Hi 10 @ 10% vol 36m",
            target_vol=0.10,
            lookback=36,
        )
        vt_mom_8_24, _ = vol_scale_monthly(raw_hi_hi, target_ann=0.08, lookback=24)
        vt_st_8_24, _ = vol_scale_monthly(st_rev, target_ann=0.08, lookback=24)
        vt_mom_10_36, _ = vol_scale_monthly(raw_hi_hi, target_ann=0.10, lookback=36)
        vt_st_10_36, _ = vol_scale_monthly(st_rev, target_ann=0.10, lookback=36)

        structural = {
            "struct_mom_8_24": (vt_mom_8_24, "Mom hi−hi10 vol-target 8% / 24m (no overlay)"),
            "struct_blend50_8_24": (
                0.5 * vt_mom_8_24 + 0.5 * vt_st_8_24,
                "50% mom + 50% ST rev @ 8% vol 24m (no overlay)",
            ),
            "struct_blend50_10_36": (
                0.5 * vt_mom_10_36 + 0.5 * vt_st_10_36,
                "50% mom + 50% ST rev @ 10% vol 36m (no overlay)",
            ),
        }
        for key, (ser, desc) in structural.items():
            strategy_rets[key] = ser
            strategy_meta_extra[key] = {"description": desc}

        overlays = {
            "low_dd_mom_8_24_dd_cut": (
                drawdown_exposure_overlay(vt_mom_8_24, trigger_frac=-0.08, reduced_exposure=0.0),
                "Mom 8%/24m; flat if strategy DD < −8%",
            ),
            "low_dd_blend75_dd_cut": (
                drawdown_exposure_overlay(
                    0.75 * vt_mom_10_36 + 0.25 * vt_st_10_36,
                    trigger_frac=-0.08,
                    reduced_exposure=0.0,
                ),
                "75% mom + 25% ST @ 10%; flat if DD < −8%",
            ),
            "low_dd_blend50_8_24_dd_cut": (
                drawdown_exposure_overlay(
                    0.5 * vt_mom_8_24 + 0.5 * vt_st_8_24,
                    trigger_frac=-0.08,
                    reduced_exposure=0.0,
                ),
                "50% mom + 50% ST @ 8%/24m; flat if DD < −8%",
            ),
        }
        for key, (ser, desc) in overlays.items():
            strategy_rets[key] = ser
            strategy_meta_extra[key] = {
                "description": desc,
                "dd_overlay_trigger_pct": -8.0,
                "dd_overlay_exposure": 0.0,
            }

    if preset == "all":
        for key, (_label, long_col, short_col) in STRATEGIES.items():
            strategy_rets[key] = _leg_spread(mom_dec, var_dec, long_col, short_col)

    rows: list[pd.DataFrame] = []
    meta_strats = {}
    for key, series in strategy_rets.items():
        st = _stats(series, capital)
        desc = strategy_meta_extra.get(key, {}).get("description") or STRATEGIES.get(key, (key,))[0]
        meta_strats[key] = {
            **st,
            "description": desc,
            **_invested_stats(series),
            **strategy_meta_extra.get(key, {}),
        }
        rows.append(_equity_frame(
            series,
            scales_map.get(key),
            raw_map.get(key),
            capital,
            key,
        ))

    out_prefix.parent.mkdir(parents=True, exist_ok=True)
    monthly_all = pd.concat(rows, ignore_index=True)
    monthly_all.to_csv(f"{out_prefix}_monthly.csv", index=False)

    wide = pd.DataFrame(strategy_rets)
    wide.index = wide.index.strftime("%Y-%m-%d")
    wide.to_csv(f"{out_prefix}_monthly_returns_wide.csv", index=False)

    cmd = (
        f"cd {_REPO} && PYTHONUNBUFFERED=1 .venv/bin/python "
        f"RenTech/strategy_stack/run_french_mom_vol_decile_backtest.py "
        f"--preset {preset} --target-vol {target_vol} --vol-lookback {vol_lookback} "
        f"--capital {int(capital)} --weighting {weighting}"
    )
    if start:
        cmd += f" --start {start}"
    if end:
        cmd += f" --end {end}"
    cmd += f" --out-prefix {out_prefix}"

    if preset == "invested":
        canonical_key = "invested_mom_10_lb24"
    elif preset == "low_dd":
        canonical_key = "low_dd_mom_8_24_dd_cut"
    elif preset == "optimized":
        canonical_key = f"vol_blend_75_mom_{int(target_vol * 100)}_strev"
    else:
        canonical_key = f"vol_target_mom_hi_{int(target_vol * 100)}"
    meta = {
        "preset": preset,
        "canonical_strategy": canonical_key,
        "source": "Ken French Data Library (portfolio returns, not stock intersection)",
        "momentum_file": str(FRENCH_DIR / "10_Portfolios_Prior_12_2.csv"),
        "variance_file": str(FRENCH_DIR / "Portfolios_Formed_on_VAR.csv"),
        "st_reversal_file": str(FRENCH_DIR / "10_Portfolios_Prior_1_0.csv"),
        "weighting": weighting,
        "target_vol_ann": target_vol,
        "vol_lookback_months": vol_lookback,
        "note_approximation": (
            "Spread is Hi PRIOR minus Hi 10 variance decile returns, "
            "not the intersection portfolio of stocks in both legs."
        ),
        "corr_lo_mom_decile_hi_var_decile": round(rho_lo_mom_hi_vol, 4),
        "strategies": meta_strats,
        "command": cmd,
    }
    yearly_path = f"{out_prefix}_{canonical_key}_yearly.csv"
    _yearly_chained(strategy_rets[canonical_key], capital).to_csv(yearly_path, index=False)
    meta["yearly_csv"] = str(yearly_path)

    with open(f"{out_prefix}_meta.json", "w") as f:
        json.dump(meta, f, indent=2)

    lines = [
        "=== French momentum / variance decile backtest ===",
        f"Preset     : {preset}",
        f"Window     : {meta_strats[canonical_key]['start']} → {meta_strats[canonical_key]['end']}",
        f"Months     : {meta_strats[canonical_key]['n_months']}",
        f"Weighting  : {weighting}",
        f"ρ(Lo PRIOR, Hi 10 var) : {rho_lo_mom_hi_vol:.4f}",
        "",
        f"{'Strategy':<32} {'Sharpe':>8} {'CAGR%':>8} {'MaxDD%':>8} {'Vol%':>8}",
        "-" * 72,
    ]
    order = list(strategy_rets.keys())
    if preset == "invested":
        order = [
            canonical_key,
            "invested_mom_8_24",
            "invested_blend50_8_24",
            "invested_blend75_10_strev",
        ]
    elif preset == "low_dd":
        order = [
            canonical_key,
            "low_dd_blend75_dd_cut",
            "low_dd_blend50_8_24_dd_cut",
            "struct_mom_8_24",
            "struct_blend50_8_24",
            "struct_blend50_10_36",
            "ref_vol_target_mom_hi_10",
        ]
    elif preset == "all":
        for k in STRATEGIES:
            if k in order:
                order.remove(k)
        order = [canonical_key] + [k for k in STRATEGIES] + [
            k for k in order if k != canonical_key and k not in STRATEGIES
        ]
    for key in order:
        if key not in meta_strats:
            continue
        s = meta_strats[key]
        lines.append(
            f"{key:<32} {s['sharpe_monthly']:>8.3f} {s['cagr_pct']:>8.2f} "
            f"{s['max_drawdown_pct']:>8.2f} {s.get('vol_ann_pct', 0):>8.2f}"
        )
    if preset == "low_dd" and "ref_vol_target_mom_hi_10" in meta_strats:
        ref = meta_strats["ref_vol_target_mom_hi_10"]
        can = meta_strats[canonical_key]
        lines.extend([
            "",
            "Vs reference (10% vol / 36m mom spread):",
            f"  MaxDD {ref['max_drawdown_pct']:.2f}% → {can['max_drawdown_pct']:.2f}%  "
            f"({can['max_drawdown_pct'] - ref['max_drawdown_pct']:.2f} pp)",
            f"  CAGR {ref['cagr_pct']:.2f}% → {can['cagr_pct']:.2f}%",
        ])
    elif "mom_ls" in meta_strats and canonical_key in meta_strats:
        lines.extend([
            "",
            "Vs classic momentum:",
            f"  DD improvement (pp): "
            f"{meta_strats[canonical_key]['max_drawdown_pct'] - meta_strats['mom_ls']['max_drawdown_pct']:.2f}",
        ])
    lines.extend(["", "Command:", f"  {cmd}"])
    metrics_txt = "\n".join(lines) + "\n"
    with open(f"{out_prefix}_metrics.txt", "w") as f:
        f.write(metrics_txt)

    if verbose:
        print(metrics_txt)
        print(f"Yearly returns: {yearly_path}")

    return meta


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--start", default="1963-07-01", help="First month (variance data starts 1963-07)")
    ap.add_argument("--end", default=None)
    ap.add_argument("--capital", type=float, default=100_000.0)
    ap.add_argument("--weighting", choices=("value", "equal"), default="value")
    ap.add_argument(
        "--preset",
        choices=("all", "vol_target_mom_hi", "optimized", "low_dd", "invested"),
        default="invested",
        help="invested = usually in market; low_dd = DD cash cut; optimized; vol_target_mom_hi; all",
    )
    ap.add_argument("--target-vol", type=float, default=0.10, help="Annual vol target for scaling")
    ap.add_argument("--vol-lookback", type=int, default=36, help="Trailing months for vol estimate")
    ap.add_argument(
        "--out-prefix",
        type=Path,
        default=LOGS / "french_mom_vol_invested",
    )
    ap.add_argument("-q", "--quiet", action="store_true")
    args = ap.parse_args()

    run_backtest(
        start=args.start,
        end=args.end,
        capital=args.capital,
        weighting=args.weighting,
        preset=args.preset,
        target_vol=args.target_vol,
        vol_lookback=args.vol_lookback,
        out_prefix=args.out_prefix,
        verbose=not args.quiet,
    )


if __name__ == "__main__":
    main()