"""
Does a factor combo that did well recently keep doing well?

For every pair of price factors (equal-weight rank composite) compute a cheap
daily proxy: the equal-weight return of its top-N members on the previous
close's scores, minus the member average. Then, each calendar year, rank the
pairs by their trailing 1/3/5-year proxy Sharpe and look at how the chosen
pairs do the NEXT year, against the average pair and a fixed pick.

    PYTHONPATH=. uv run python -m research.lab.persistence
"""

from __future__ import annotations

import itertools
import os
import pickle
import time
import warnings

import numpy as np
import pandas as pd
from scipy import stats as sps

from research.lab.factor_scan import RESULTS, load
from research.lab.rotation import factor_ranks

warnings.filterwarnings("ignore")


def top_proxy(p, score: pd.DataFrame, top_n: int = 20) -> pd.Series:
    """Daily excess return of the top-N (previous close scores) over members."""
    S = score.shift(1).values
    r = p.close.pct_change(fill_method=None).values
    mem = p.member.shift(1, fill_value=False).values
    S = np.where(mem & ~np.isnan(r), S, np.nan)
    out = np.full(len(S), np.nan)
    base = np.nanmean(np.where(mem, r, np.nan), axis=1)
    for t in range(len(S)):
        row = S[t]
        ok = np.where(~np.isnan(row))[0]
        if len(ok) < top_n * 3:
            continue
        top = ok[np.argpartition(-row[ok], top_n)[:top_n]]
        out[t] = np.nanmean(r[t, top]) - base[t]
    return pd.Series(out, index=p.dates)


def main() -> None:
    p, _ = load("us")
    ranks = factor_ranks(p, "us")
    names = sorted(ranks)
    path = os.path.join(RESULTS, "pair_proxies_us.pkl")
    if os.path.exists(path):
        with open(path, "rb") as fh:
            P = pickle.load(fh)
    else:
        t = time.time()
        P = {}
        for a, b in itertools.combinations(names, 2):
            s = (ranks[a] + ranks[b]) / 2
            P[f"{a}+{b}"] = top_proxy(p, s)
        for a in names:
            P[a] = top_proxy(p, ranks[a])
        P = pd.DataFrame(P)
        print(f"{P.shape[1]} proxies [{time.time() - t:.0f}s]", flush=True)
        with open(path, "wb") as fh:
            pickle.dump(P, fh)
    P = P.loc["2006":]
    yearly = P.groupby(P.index.year).apply(lambda x: x.mean() / x.std() * np.sqrt(252))   # excess Sharpe (IR)
    yret = P.groupby(P.index.year).sum()                                                    # ~excess return
    years = [y for y in yearly.index if y >= 2008 and y <= 2025]
    print("\nRank correlation of combo IR between consecutive years (all pairs):")
    rho = [sps.spearmanr(yearly.loc[y - 1], yearly.loc[y], nan_policy="omit").statistic for y in years]
    print("  " + " ".join(f"{y}:{r:+.2f}" for y, r in zip(years, rho)), f"| mean {np.nanmean(rho):+.3f}")
    for look in (1, 3, 5):
        rows = []
        for y in range(2006 + look, 2027):
            hist = P.loc[str(y - look):str(y - 1)]
            ir = hist.mean() / hist.std()
            nxt = yret.loc[y] if y in yret.index else None
            if nxt is None:
                continue
            best = ir.nlargest(10).index
            worst = ir.nsmallest(10).index
            rows.append({"year": y, "top10": nxt[best].mean(), "bottom10": nxt[worst].mean(),
                         "all": nxt.mean(), "best1": nxt[ir.idxmax()], "best1_name": ir.idxmax()})
        df = pd.DataFrame(rows).set_index("year")
        print(f"\nlookback {look}y: next-year excess return of the pairs ranked by trailing IR")
        print(f"  mean  top10 {df['top10'].mean()*100:+.2f}%  bottom10 {df['bottom10'].mean()*100:+.2f}%  "
              f"all {df['all'].mean()*100:+.2f}%  best1 {df['best1'].mean()*100:+.2f}%  "
              f"top10 beats all in {(df['top10'] > df['all']).mean()*100:.0f}% of years")
        if look == 1:
            print(df.assign(**{c: (df[c] * 100).round(1) for c in ["top10", "bottom10", "all", "best1"]}).to_string())


if __name__ == "__main__":
    main()