"""
Theme concentration in the momentum sleeve. In the 2026-04..09 replay the
live overnight_recovery_top10 held about 92% in AI-hardware names; this
checks whether capping names per industry changes its return or drawdown
over 2006-2026.

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

A cap of k per group keeps a stock eligible only while it ranks in the top k
of its group by the same score; the usual top-10 / buffer-2 / every-10-days
selection then runs on the eligible names. Groups: GICS sector (current
classification, 11 groups) and 3-digit SIC (sec_sic.csv, ~150 groups).
Reported: alpha vs QQQ per half, MaxDD, worst 3-month return, and the mean
largest group weight.
"""

from __future__ import annotations

import os
import warnings

import numpy as np
import pandas as pd

from research.lab.alpha import alpha_stats, bench_returns
from research.lab.factor_scan import RESULTS, load
from research.lab.rotation import factor_ranks
from research.lab.us_sweep import _blend, run

warnings.filterwarnings("ignore")
LONG = [("2006-09-27", "2026-09-24")]


def capped(score: pd.DataFrame, groups: pd.Series, k: int) -> pd.DataFrame:
    g = groups.reindex(score.columns)
    out = score.copy()
    for name in g.dropna().unique():
        cols = list(g.index[g == name])
        rk = score[cols].rank(axis=1, ascending=False)
        out[cols] = score[cols].where(rk <= k)
    return out


def top_group_weight(score: pd.DataFrame, groups: pd.Series, n: int = 10) -> float:
    rk = score.rank(axis=1, ascending=False) <= n
    g = groups.reindex(score.columns)
    share = [rk.loc[:, g == x].sum(axis=1) for x in g.dropna().unique()]
    m = pd.concat(share, axis=1).max(axis=1) / n
    return float(m[rk.sum(axis=1) >= n].mean())


def main() -> None:
    p, _ = load("us")
    pr = factor_ranks(p, "us", ["overnight_mom", "rec_126"])
    mom = _blend(pr["overnight_mom"], pr["rec_126"]).where(p.member)
    sic = pd.read_csv(os.path.join("data", "research", "sec_sic.csv")).drop_duplicates("ticker").set_index("ticker")["sic"]
    sic3 = (pd.to_numeric(sic, errors="coerce") // 10).dropna().astype(int).astype(str)
    groups = {"gics": p.extra["sector"], "sic3": sic3}
    b = bench_returns(p.dates)
    variants = [("uncapped", None, None), ("gics <=4", "gics", 4), ("gics <=3", "gics", 3), ("gics <=2", "gics", 2),
                ("sic3 <=3", "sic3", 3), ("sic3 <=2", "sic3", 2), ("sic3 <=1", "sic3", 1)]
    rows, R = [], {}
    for name, gk, k in variants:
        s = mom if gk is None else capped(mom, groups[gk], k)
        r = run(p, s, LONG, top_n=10, buffer=2.0, every=10)["_rets"]
        e = (1 + r).cumprod()
        q = (1 + r).rolling(63).apply(np.prod, raw=True) - 1
        row = {"variant": name, "mdd": float((e / e.cummax() - 1).min()), "worst_3m": float(q.min()),
               "vol": float(r.std() * np.sqrt(252)),
               "max_gics_w": top_group_weight(s.loc["2006-09-27":], groups["gics"]),
               "max_sic3_w": top_group_weight(s.loc["2006-09-27":], groups["sic3"])}
        for h, x in (("all", r), ("design", r.loc[:"2016"]), ("holdout", r.loc["2017":])):
            a = alpha_stats(x, b)
            row.update({f"{h}_alpha": a["alpha"], f"{h}_t": a["t"], f"{h}_cagr": a["cagr"]})
        row["sharpe"] = float(r.mean() / r.std() * np.sqrt(252))
        rows.append(row)
        R[name] = r
        print(f"{name:10s} alpha {row['all_alpha']*100:5.1f}% (t {row['all_t']:.2f}) | 06-16 {row['design_alpha']*100:5.1f}% "
              f"t {row['design_t']:.2f} | 17-26 {row['holdout_alpha']*100:5.1f}% t {row['holdout_t']:.2f} | "
              f"cagr {row['all_cagr']*100:5.1f}% sharpe {row['sharpe']:.2f} vol {row['vol']*100:4.1f}% "
              f"mdd {row['mdd']*100:4.0f}% worst3m {row['worst_3m']*100:4.0f}% | max sector w {row['max_gics_w']*100:3.0f}% "
              f"max sic3 w {row['max_sic3_w']*100:3.0f}%", flush=True)
    pd.DataFrame(rows).to_csv(os.path.join(RESULTS, "concentration.csv"), index=False)
    pd.DataFrame(R).to_pickle(os.path.join(RESULTS, "concentration_returns.pkl"))


if __name__ == "__main__":
    main()