quant-trading/research/lab/review_live_us.py

"""
Review of the US replay books over 2026-04-08 .. 2026-09-28: what drove the
market, where each strategy's return came from, and how good its entries and
exits were.

    PYTHONPATH=. uv run python -m research.lab.review_live_us

Inputs are the live data files (data/us*.csv, data/etfs.csv, the original
book's data/trader_us_*.json). The books are point-in-time replays, not live
trading (see project memory); every decision in them used only data up to
its date. Books trade at the close, so the weights recorded on day d earn
the close-to-close return of day d+1.

AI_CHAIN groups S&P 500 members into stages of the AI build-out. The grouping
was written in 2026 with hindsight and is used here to describe the period,
never as a trading signal.
"""

from __future__ import annotations

import glob
import json
import os
import warnings

import numpy as np
import pandas as pd

from universe_history import membership_mask

warnings.filterwarnings("ignore")
START, END = "2026-04-08", "2026-09-28"
RESULTS = os.path.join("data", "research", "results")

AI_CHAIN = {
    "hyperscaler": ["MSFT", "GOOGL", "AMZN", "META", "ORCL"],
    "ai_software": ["PLTR", "NOW", "CRM", "ADBE", "APP", "INTU", "DDOG", "CRWD", "SNOW", "WDAY"],
    "servers_network": ["DELL", "SMCI", "HPE", "ANET", "CSCO", "JBL", "FLEX"],
    "compute_chips": ["NVDA", "AMD", "AVGO", "MRVL", "INTC", "QCOM"],
    "memory_storage": ["MU", "SNDK", "WDC", "STX", "NTAP", "PSTG"],
    "optical_components": ["LITE", "COHR", "CIEN", "GLW", "APH", "TEL"],
    "semi_equipment": ["AMAT", "LRCX", "KLAC", "TER", "Q"],
    "power_datacenter": ["VST", "CEG", "NRG", "GEV", "ETN", "VRT", "BE", "PWR", "FIX", "EME", "EQIX", "DLR"],
    "analog_semis": ["TXN", "ADI", "MCHP", "ON", "NXPI", "MPWR"],
}
FOCUS = ["overnight_recovery_top10", "us_momentum_qqq_trend", "us_three_sleeve", "lowrisk_top20",
         "fc_up_cap_mom_gap_weekly", "fc_rec_mfilt_deep_upvol_monthly", "recovery_mom_top10",
         "momentum", "trend_rider_v7_vt24"]


def load_prices():
    c = pd.read_csv("data/us.csv", index_col=0, parse_dates=True)
    e = pd.read_csv("data/etfs.csv", index_col=0, parse_dates=True)
    c = c.join(e[[x for x in e.columns if x not in c.columns]], how="left")
    c = c.sort_index()
    etf = [x for x in e.columns if x in c.columns]
    c[etf] = c[etf].ffill(limit=3)
    return c


def load_books() -> dict[str, dict]:
    out = {}
    for f in sorted(glob.glob("data/trader_us_*.json")):
        d = json.load(open(f))
        if d.get("daily_log"):
            out[d["strategy"]] = d
    return out


def book_weights(book: dict, cols) -> tuple[pd.DataFrame, pd.Series]:
    """Weights (value / equity) after each day's close, and daily equity."""
    rows, eq = {}, {}
    for r in book["daily_log"]:
        e = r["equity"]
        eq[r["date"]] = e
        rows[r["date"]] = {t: h["value"] / e for t, h in r["holdings"].items()}
    w = pd.DataFrame.from_dict(rows, orient="index").fillna(0.0)
    w.index = pd.to_datetime(w.index)
    eq = pd.Series(eq); eq.index = pd.to_datetime(eq.index)
    w, eq = w.sort_index(), eq.sort_index()
    return w[~w.index.duplicated(keep="last")], eq[~eq.index.duplicated(keep="last")]


def stage_of() -> dict[str, str]:
    return {t: s for s, ts in AI_CHAIN.items() for t in ts}


def market_view(C: pd.DataFrame, members: list[str], sectors: pd.Series) -> dict:
    r = C.loc[START:END, members].ffill()
    tot = (r.iloc[-1] / r.iloc[0] - 1).dropna()
    st = stage_of()
    df = pd.DataFrame({"ret": tot, "sector": sectors.reindex(tot.index),
                       "stage": [st.get(t, "") for t in tot.index]})
    out = {"n": len(df), "ew": float(df["ret"].mean()), "median": float(df["ret"].median()),
           "spy": float(C.loc[END, "SPY"] / C.loc[START, "SPY"] - 1),
           "qqq": float(C.loc[END, "QQQ"] / C.loc[START, "QQQ"] - 1),
           "smh": float(C.loc[END, "SMH"] / C.loc[START, "SMH"] - 1)}
    out["by_sector"] = df.groupby("sector")["ret"].agg(["mean", "median", "count"]).sort_values("mean", ascending=False)
    out["by_stage"] = df[df["stage"] != ""].groupby("stage")["ret"].agg(["mean", "median", "count"]).sort_values("mean", ascending=False)
    top = df.sort_values("ret", ascending=False).head(40)
    out["top40"] = top
    out["top40_ai_share"] = float((top["stage"].isin([s for s in AI_CHAIN if s != "analog_semis"])).mean())
    return out


def attribution(w: pd.DataFrame, R: pd.DataFrame, bench_members: pd.DataFrame, groups: pd.Series) -> dict:
    """Daily Brinson against the equal-weight PIT S&P 500: group allocation and
    selection within groups, summed arithmetically over the period."""
    dates = w.index.intersection(R.index)
    alloc = sel = act = 0.0
    per_group = {}
    for d0, d1 in zip(dates[:-1], dates[1:]):
        wp = w.loc[d0]
        wp = wp[wp > 0]
        r1 = R.loc[d1]
        mem = bench_members.loc[d0]
        bm = [t for t in mem.index[mem] if np.isfinite(r1.get(t, np.nan))]
        g_b = groups.reindex(bm).fillna("Other")
        rb = r1[bm]
        Rb = rb.mean()
        g_p = groups.reindex(wp.index).fillna("Other")
        rp = r1.reindex(wp.index).fillna(0.0)
        for g in set(g_b) | set(g_p):
            wb_g = float((g_b == g).mean())
            rb_g = float(rb[g_b == g].mean()) if (g_b == g).any() else Rb
            wp_g = float(wp[g_p == g].sum())
            rp_g = float((wp[g_p == g] * rp[g_p == g]).sum() / wp_g) if wp_g > 0 else rb_g
            a = (wp_g - wb_g) * (rb_g - Rb)
            s = wp_g * (rp_g - rb_g)
            alloc += a; sel += s
            pg = per_group.setdefault(g, [0.0, 0.0])
            pg[0] += a; pg[1] += s
        cash = 1 - wp.sum()
        act += float((wp * rp).sum()) - Rb      # cash earns 0 in the books
    pg = pd.DataFrame(per_group, index=["allocation", "selection"]).T
    pg["total"] = pg.sum(axis=1)
    return {"allocation": alloc, "selection": sel, "active_sum": act, "by_group": pg.sort_values("total")}


def trade_quality(book: dict, C: pd.DataFrame, ew: pd.Series, horizons=(10, 21)) -> pd.DataFrame:
    """Forward return after each full exit / entry, minus the equal-weight index."""
    rows = []
    idx = C.index
    for t in book["trade_log"]:
        d = pd.Timestamp(t["date"])
        if t["ticker"] not in C.columns or d not in idx:
            continue
        i = idx.get_loc(d)
        row = {"date": d, "ticker": t["ticker"], "action": t["action"]}
        for h in horizons:
            if i + h < len(idx):
                p0, p1 = C[t["ticker"]].iloc[i], C[t["ticker"]].iloc[i + h]
                b = ew.iloc[i + 1:i + h + 1]
                row[f"fwd{h}"] = p1 / p0 - 1 - ((1 + b).prod() - 1)
        rows.append(row)
    return pd.DataFrame(rows)


def exits_entries(w: pd.DataFrame) -> tuple[list, list]:
    held = w > 1e-6
    prev = held.shift(1, fill_value=False)
    entries = [(d, t) for d, t in zip(*np.where(held & ~prev)) if d > 0]
    exits = [(d, t) for d, t in zip(*np.where(~held & prev))]
    return ([(w.index[d], w.columns[t]) for d, t in entries],
            [(w.index[d], w.columns[t]) for d, t in exits])


def winner_capture(w: pd.DataFrame, C: pd.DataFrame, winners: list[str]) -> pd.DataFrame:
    """For each period winner: days held, first entry, and the log-return earned
    while held as a share of its total log-return over the period."""
    R = np.log(C.loc[START:END].ffill()).diff()
    rows = []
    for t in winners:
        tot = float(R[t].sum())
        if t in w.columns:
            h = (w[t] > 1e-6).reindex(R.index).fillna(False).shift(1, fill_value=False)
            got = float(R[t][h].sum())
            first = w.index[w[t] > 1e-6][0] if (w[t] > 1e-6).any() else None
            days = int(h.sum())
        else:
            got, first, days = 0.0, None, 0
        rows.append({"ticker": t, "total_logret": tot, "captured_logret": got,
                     "capture": got / tot if tot else np.nan, "days_held": days, "first_entry": first})
    return pd.DataFrame(rows)


def stage_timing(C: pd.DataFrame) -> pd.DataFrame:
    """Equal-weight stage indices from 2025-01; the date each one first reached
    25/50/75% of its log gain up to END."""
    rows = {}
    px = C.loc["2025-01-01":END].ffill()
    for s, ts in AI_CHAIN.items():
        ts = [t for t in ts if t in px.columns]
        r = np.log(px[ts]).diff().mean(axis=1).fillna(0.0)
        cum = r.cumsum()
        tot = cum.iloc[-1]
        rows[s] = {"logret_2025_01_to_end": float(tot)}
        for q in (0.25, 0.5, 0.75):
            hit = cum.index[cum >= q * tot]
            rows[s][f"t{int(q*100)}"] = hit[0].date() if len(hit) and tot > 0 else None
    return pd.DataFrame(rows).T


def main() -> None:
    os.makedirs(RESULTS, exist_ok=True)
    C = load_prices()
    sec = pd.read_csv("data/us_sectors.csv").set_index("ticker")
    from research.lab.factor_scan import load
    p, _ = load("us")
    sectors = p.extra["sector"].combine_first(sec["GICS Sector"])
    dates = C.loc[START:END].index
    M = membership_mask(dates, tickers=[c for c in C.columns if c in sectors.index])
    members0 = list(M.columns[M.iloc[0]])
    R = C.pct_change(fill_method=None)
    ew = R[members0].loc[dates].mean(axis=1)

    mv = market_view(C, members0, sectors)
    print(f"period {START}..{END}: SPY {mv['spy']*100:.1f}% QQQ {mv['qqq']*100:.1f}% SMH {mv['smh']*100:.1f}% "
          f"EW members {mv['ew']*100:.1f}% (median {mv['median']*100:.1f}%, n {mv['n']})")
    print("\nby GICS sector (equal-weight member return):")
    print((mv["by_sector"].assign(mean=lambda x: x["mean"] * 100, median=lambda x: x["median"] * 100)).round(1).to_string())
    print("\nby AI-chain stage (hindsight grouping):")
    print((mv["by_stage"].assign(mean=lambda x: x["mean"] * 100, median=lambda x: x["median"] * 100)).round(1).to_string())
    print(f"\ntop 40 winners: {mv['top40_ai_share']*100:.0f}% are AI-chain names")
    print((mv["top40"].assign(ret=lambda x: x["ret"] * 100)).round(1).to_string())
    print("\nstage timing (EW index from 2025-01; date first reaching 25/50/75% of the gain):")
    timing = stage_timing(C)
    print(timing.to_string())

    books = load_books()
    st = stage_of()
    grp_stage = pd.Series({t: st.get(t, "not_ai") for t in C.columns})
    winners = list(mv["top40"].index)
    summary, attrib, caps, tq_all = [], {}, {}, {}
    for name, b in books.items():
        w, eq = book_weights(b, C.columns)
        w = w.loc[START:END]
        ret = float(eq.loc[:END].iloc[-1] / b["initial_capital"] - 1)
        sd = str(eq.index[0].date())
        spy = float(C.loc[END, "SPY"] / C.loc[sd:, "SPY"].iloc[0] - 1) if sd else np.nan
        stk = w[[c for c in w.columns if c in sectors.index]]
        ai = stk[[c for c in stk.columns if st.get(c, "") and st.get(c) != "analog_semis"]].sum(axis=1)
        tq = trade_quality(b, C, R[members0].mean(axis=1))
        row = {"strategy": name, "start": sd, "ret": ret, "spy_same_span": spy,
               "stock_w": float(stk.sum(axis=1).mean()), "ai_w": float(ai.mean()),
               "n_trades": len(b["trade_log"])}
        if len(tq):
            s = tq[tq["action"] == "SELL"]; bu = tq[tq["action"] == "BUY"]
            row.update({"sell_fwd21": s["fwd21"].mean() if "fwd21" in s else np.nan,
                        "buy_fwd21": bu["fwd21"].mean() if "fwd21" in bu else np.nan})
        cap = winner_capture(w, C, winners)
        row["winner_capture"] = float(cap["captured_logret"].sum() / cap["total_logret"].sum())
        row["winners_held"] = int((cap["days_held"] > 0).sum())
        summary.append(row)
        if name in FOCUS and sd <= START:
            attrib[name] = {"sector": attribution(w, R, M, sectors), "ai_stage": attribution(w, R, M, grp_stage)}
            caps[name] = cap
            tq_all[name] = tq
    S = pd.DataFrame(summary).set_index("strategy").sort_values("ret", ascending=False)
    print("\nall books (ret since own start; ai_w = mean weight in AI-chain names; winner_capture = share of the top-40 "
          "winners' log gain earned; sell/buy_fwd21 = 21-day excess return after the trade):")
    print(S.assign(**{c: S[c] * 100 for c in ["ret", "spy_same_span", "stock_w", "ai_w", "sell_fwd21", "buy_fwd21", "winner_capture"]})
          .round(1).to_string())
    full = S[S["start"] <= START]
    print(f"\nbooks from {START}: corr(ret, ai_w) {full['ret'].corr(full['ai_w']):.2f}, "
          f"corr(ret, winner_capture) {full['ret'].corr(full['winner_capture']):.2f}")

    for name, a in attrib.items():
        print(f"\n== {name}: Brinson vs EW S&P 500 (sum of daily terms, %)")
        for k in ("sector", "ai_stage"):
            x = a[k]
            print(f"  by {k}: allocation {x['allocation']*100:6.1f}  selection {x['selection']*100:6.1f}  "
                  f"active {x['active_sum']*100:6.1f}")
            print("   ", ", ".join(f"{g} {v*100:+.1f}" for g, v in x["by_group"]["total"].items() if abs(v) > 0.01))
        cap = caps[name]
        held = cap[cap["days_held"] > 0].sort_values("first_entry")
        print("  top-40 winners held (first entry, capture):",
              ", ".join(f"{r.ticker} {str(r.first_entry.date())[5:]} {r.capture*100:.0f}%" for r in held.itertuples()))
        miss = cap[cap["days_held"] == 0].sort_values("total_logret", ascending=False)
        print("  never held:", ", ".join(miss["ticker"].head(20)))
        tq = tq_all[name]
        if len(tq):
            s = tq[tq["action"] == "SELL"]
            print(f"  exits: {len(s)}, 21d excess after exit mean {s['fwd21'].mean()*100:+.1f}% "
                  f"median {s['fwd21'].median()*100:+.1f}%, share >+10%: {(s['fwd21'] > .10).mean()*100:.0f}%")

    out = {"market": {k: v for k, v in mv.items() if not isinstance(v, pd.DataFrame)},
           "summary": S.reset_index().to_dict(orient="records"),
           "by_sector": mv["by_sector"].reset_index().to_dict(orient="records"),
           "by_stage": mv["by_stage"].reset_index().to_dict(orient="records"),
           "top40": mv["top40"].reset_index().rename(columns={"index": "ticker"}).to_dict(orient="records"),
           "timing": timing.reset_index().astype(str).to_dict(orient="records"),
           "attribution": {n: {k: {"allocation": v["allocation"], "selection": v["selection"],
                                   "active": v["active_sum"], "by_group": v["by_group"]["total"].to_dict()}
                               for k, v in a.items()} for n, a in attrib.items()},
           "capture": {n: c.astype({"first_entry": str}).to_dict(orient="records") for n, c in caps.items()}}
    with open(os.path.join(RESULTS, "review_live_us.json"), "w") as fh:
        json.dump(out, fh, indent=1, default=str)
    pd.to_pickle({"trades": tq_all}, os.path.join(RESULTS, "review_live_us_trades.pkl"))


if __name__ == "__main__":
    main()