"""
Industry momentum (Moskowitz & Grinblatt 1999) on S&P 500 members, with
industries from SEC SIC codes (data/research/sec_sic.csv, written by
research/lab/fetch_lazy_prices.py) grouped to 2-digit SIC.

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

Industry return = equal-weight mean of its PIT members' daily returns (groups
with fewer than 3 members that day are left out). Scores, known at each close:
  ind_mom_6_1    industry return over months 6..1 (skip the last month)
  ind_mom_12_1   same over 12..1
Uses: (a) a sleeve holding the top-20 stocks by their industry's momentum,
(b) the momentum sleeve restricted to stocks in the top half of industries.
Alpha against QQQ, 2006-2016 vs 2017-2026.
"""

from __future__ import annotations

import dataclasses
import os
import warnings

import numpy as np
import pandas as pd

from research.lab.factors import _ret

warnings.filterwarnings("ignore")
SIC = os.path.join("data", "research", "sec_sic.csv")


def industries(p) -> pd.Series:
    s = pd.read_csv(SIC).drop_duplicates("ticker").set_index("ticker")["sic"]
    s = pd.to_numeric(s, errors="coerce").dropna().astype(int) // 100
    return s.reindex(p.close.columns)


def industry_momentum(p, lookback: int, skip: int = 21) -> pd.DataFrame:
    """Each stock gets its industry's cumulative return over [t-lookback, t-skip]."""
    ind = industries(p)
    r = _ret(p).where(p.member)
    out = pd.DataFrame(np.nan, index=p.dates, columns=p.close.columns)
    for g in ind.dropna().unique():
        cols = list(ind.index[ind == g])
        n = r[cols].notna().sum(axis=1)
        gr = r[cols].mean(axis=1).where(n >= 3)
        cum = np.log1p(gr.fillna(0.0)).rolling(lookback - skip, min_periods=(lookback - skip) // 2).sum().shift(skip)
        out[cols] = np.repeat(cum.where(n >= 3).to_numpy()[:, None], len(cols), axis=1)
    return out.where(p.member)


def main() -> None:
    from research.lab.alpha import alpha_stats, bench_returns, fmt_alpha
    from research.lab.factor_scan import load
    from research.lab.rotation import factor_ranks
    from research.lab.us_sweep import _blend, run

    p, _ = load("us")
    ind = industries(p)
    print(f"{ind.notna().sum()}/{len(ind)} tickers with SIC; {ind.nunique()} 2-digit groups")
    b = bench_returns(p.dates)
    pr = factor_ranks(p, "us")
    mom = _blend(pr["overnight_mom"], pr["rec_126"])
    span = [("2006-09-27", "2026-09-24")]

    def show(name, r):
        a, d, h = alpha_stats(r, b), alpha_stats(r.loc[:"2016"], b), alpha_stats(r.loc["2017":], b)
        print(f"{fmt_alpha(name, a)} | 06-16 a {d['alpha']*100:5.1f}% t {d['t']:5.2f} | 17-26 a {h['alpha']*100:5.1f}% t {h['t']:5.2f}", flush=True)

    for lb in (126, 252):
        im = industry_momentum(p, lb)
        rk = im.rank(axis=1, pct=True)
        show(f"industry mom {lb}d top20 e21", run(p, rk, span, top_n=20, buffer=2.0, every=21)["_rets"])
        tie = rk + 0.001 * pr["mom_12_1"].fillna(0)        # stocks of the best industries, by own momentum
        show(f"industry mom {lb}d (+stock tiebreak) top20", run(p, tie.where(p.member), span, top_n=20, buffer=2.0, every=21)["_rets"])
        keep = rk >= 0.5
        pm = dataclasses.replace(p, member=p.member & keep.fillna(False))
        show(f"momentum top10, top-half industries ({lb}d)", run(pm, mom.where(pm.member), span, top_n=10, buffer=2.0, every=10)["_rets"])
    show("momentum top10 (reference)", run(p, mom.where(p.member), span, top_n=10, buffer=2.0, every=10)["_rets"])


if __name__ == "__main__":
    main()