"""
Short-horizon and residual stock sleeves for the US book: do they add a
return source the three current sleeves (momentum, low risk, trend) lack?

    PYTHONPATH=. uv run --with numba python -m research.lab.short_horizon

Candidates, price/volume only (S&P 500 PIT, next-open fills, 10 bp):
  resid_mom        residual 12-1 momentum / residual vol (Blitz, Huij & Martens 2011;
                   Blitz, Hanauer & Vidojevic 2020)
  stmom_hiturn     1-month momentum among high-turnover names (Medhat & Schmeling 2022)
  rev_sector_21    industry-adjusted 1-month reversal (Da, Liu & Schaumburg 2014)
  rev_resid_21     residual 1-month reversal
  seasonality      same-calendar-month return, years 1-10 (Heston & Sadka 2008)
  short_combo      price-only subset of the Robeco short-horizon composite (Blitz et
                   al. 2023): short-term momentum for high-turnover names, residual
                   reversal for the rest, seasonality, low 1-month idiosyncratic vol

Each is run as top 20 with a rank buffer (hold while rank <= 40), at two
rebalance speeds. Reports stand-alone stats, monthly correlation with the
current sleeves (data/research/results/us_final_sleeves.pkl, from us_final),
and the 2011-2026 book with the candidate as a fourth equal sleeve.

Last, concentration: the momentum sleeve (overnight + rec_126) with 3 or 5
names instead of 10, the other thing a small account can do that a fund
cannot.
"""

from __future__ import annotations

import os
import warnings

import pandas as pd

from research.lab.factor_scan import RESULTS, load
from research.lab.rotation import factor_ranks
from research.lab.us_final import describe, line, monthly_blend
from research.lab.us_sweep import _blend, fmt, log_trials, run

warnings.filterwarnings("ignore")

SPAN = ("2011-01-03", "2026-09-24")
LONG = ("2006-09-27", "2026-09-24")


def candidates(p) -> dict[str, pd.DataFrame]:
    pr = factor_ranks(p, "us")
    hi = pr["stmom_hiturn"].notna()
    directional = pr["stmom_hiturn"].where(hi, pr["rev_resid_21"])
    return {
        "resid_mom": pr["resid_mom"],
        "stmom_hiturn": pr["stmom_hiturn"],
        "rev_sector_21": pr["rev_sector_21"],
        "rev_resid_21": pr["rev_resid_21"],
        "seasonality": pr["seasonality"],
        "short_combo": _blend(directional, pr["seasonality"], pr["ivol_20"]),
        "resid_mom+short_combo": _blend(pr["resid_mom"], directional, pr["seasonality"], pr["ivol_20"]),
    }


def main() -> None:
    p, _ = load("us")
    base = pd.read_pickle(os.path.join(RESULTS, "us_final_sleeves.pkl"))
    cur = base[["momentum", "lowrisk", "trend"]]
    rows, R = [], {}
    for name, score in candidates(p).items():
        for every in (5, 21):
            st = run(p, score.where(p.member), [LONG], top_n=20, buffer=2.0, every=every)
            key = f"{name} e{every}"
            R[key] = st["_rets"]
            d = describe(st["_rets"].loc[SPAN[0]:])
            dl = describe(st["_rets"])
            rows.append({"score": key, **{k: v for k, v in st.items() if not k.startswith("_")}})
            print(f"{line(key + ' 2011-', d)} turn {st['turn']:4.1f}x | 2006- sh {dl['sharpe']:.2f} "
                  f"mdd {dl['mdd']*100:5.1f}%", flush=True)
    log_trials(rows, "short_horizon")
    R = pd.DataFrame(R)
    R.to_pickle(os.path.join(RESULTS, "short_horizon_sleeves.pkl"))

    span = pd.concat([cur, R], axis=1).loc[SPAN[0]:SPAN[1]].fillna(0.0)
    monthly = (1 + span).resample("ME").prod() - 1
    print("\nmonthly correlation with current sleeves, 2011-2026:")
    print(monthly.corr().loc[R.columns, ["momentum", "lowrisk", "trend"]].round(2).to_string())

    three = {"momentum": 1 / 3, "lowrisk": 1 / 3, "trend": 1 / 3}
    print("\n2011-2026 book, monthly rebalanced between sleeves:")
    print(line("current 3 sleeves", describe(monthly_blend(span, three))))
    for c in R.columns:
        four = {"momentum": .25, "lowrisk": .25, "trend": .25, c: .25}
        swap = {"momentum": 1 / 3, c: 1 / 3, "trend": 1 / 3}
        print(line(f"+ {c} (4 x 1/4)", describe(monthly_blend(span, four))))
        print(line(f"  {c} replaces lowrisk", describe(monthly_blend(span, swap))))

    print("\nconcentration: overnight + rec_126, every 10 days, 2006-2026")
    pr = factor_ranks(p, "us")
    mom = _blend(pr["overnight_mom"], pr["rec_126"]).where(p.member)
    blocks = [(LONG[0], "2010-12-31"), ("2011-01-01", "2015-12-31"), ("2016-01-01", "2020-12-31"),
              ("2021-01-01", LONG[1])]
    for n in (3, 5, 10):
        for buf in (1.0, 2.0):
            st = run(p, mom, blocks, top_n=n, buffer=buf, every=10)
            print(fmt(f"top {n} buffer {buf}", st), flush=True)


if __name__ == "__main__":
    main()