"""
Does a linked company's move predict a stock that has not moved yet? Tests
of the "X already rose, its supplier Y has not priced it in" idea on PIT
S&P 500 members, with links of four kinds:

  corr_peer     the 10 members whose daily residual returns (market beta
                removed) correlated most with the stock over the last 252
                days; re-estimated each month end. A proxy for economic
                similarity (same industry, same customers, same inputs).
  lead_peer     the 10 members whose weekly residual return in week w-1
                best predicted the stock's return in week w over the last
                104 weeks (lagged cross-correlation); re-estimated each
                month end. A data-driven lead-lag network.
  10k_customer  members named in the stock's latest 10-K in a customer
                sentence (research/lab/fetch_supply_links.py), valid from
                the filing day for 15 months. Cohen & Frazzini (2008)
                customer momentum.
  10k_supplier  members named in a supplier sentence (upstream links).
  10k_any       any member named in the 10-K, either direction (a filer's
                links plus the filers that name it).

Signals, known at the close of day t:
  <link>_mom    mean return of the linked names over the last 21 trading
                days (5 for lead_peer).
  <link>_gap    that minus the stock's own return over the same days: the
                "linked names moved, this one has not" catch-up signal.

Measures (design 2006-2016, holdout 2017-2026, 10 bp, next-open fills):
  * Fama-MacBeth: next-21-day (next-5-day for lead_peer) return rank on the
    signal rank, controlling for the stock's own 1-month return, 12-1
    momentum and overnight momentum ranks; non-overlapping periods.
  * a top-20 sleeve on the signal alone: alpha against QQQ.
  * momentum top 10 (the live overnight_recovery_top10) with the top 30 by
    momentum re-sorted by the signal, vs momentum alone.
Plus an event test: after a member's earnings day (8-K item 2.02), do its
corr_peers drift in the direction of its announcement return over the next
1-21 days (under-reaction, the idea) or reverse (over-reaction, Thomas &
Zhang 2008)?

    PYTHONPATH=. uv run --with numba python -m research.lab.lead_lag [--part signals|events|chain]
"""

from __future__ import annotations

import argparse
import json
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.factors import _ret, resid_returns
from research.lab.factor_scan import RESULTS, load
from research.lab.rotation import factor_ranks

warnings.filterwarnings("ignore")
LONG = [("2006-09-27", "2026-09-24")]
DESIGN_END = "2016-12-31"
K = 10
CACHE = os.path.join("data", "research", "lead_lag_links.pkl")


# ---------------------------------------------------------------------------
# link matrices: {month_end: (tickers_idx, W)} with W[i, j] = weight of j in i's links
# ---------------------------------------------------------------------------

def month_ends(dates: pd.DatetimeIndex) -> list[pd.Timestamp]:
    s = pd.Series(dates, index=dates)
    return list(s.groupby(dates.to_period("M")).max())


def corr_peer_links(p, win: int = 252, k: int = K) -> dict:
    e = resid_returns(p).where(p.member)
    out = {}
    X = e.to_numpy()
    for me in month_ends(p.dates):
        i1 = p.dates.get_loc(me) + 1
        if i1 < win:
            continue
        blk = X[i1 - win:i1]
        mem = p.member.iloc[i1 - 1].to_numpy() & (np.isfinite(blk).sum(axis=0) >= int(win * 0.8))
        idx = np.where(mem)[0]
        Z = blk[:, idx]
        Z = (Z - np.nanmean(Z, axis=0)) / np.nanstd(Z, axis=0)
        Z = np.nan_to_num(Z)
        C = Z.T @ Z / win
        np.fill_diagonal(C, -np.inf)
        top = np.argsort(-C, axis=1)[:, :k]
        out[me] = (idx, idx[top])
    return out


def lead_peer_links(p, weeks: int = 104, k: int = K) -> dict:
    e = resid_returns(p).where(p.member).fillna(0.0)
    cum = e.cumsum().to_numpy()
    out = {}
    for me in month_ends(p.dates):
        i1 = p.dates.get_loc(me)
        if i1 < 5 * (weeks + 1) + 252:
            continue
        pts = np.arange(i1, i1 - 5 * (weeks + 1), -5)[::-1]
        W = np.diff(cum[pts], axis=0)                          # (weeks, N) weekly residual returns
        mem = p.member.iloc[i1].to_numpy() & (np.abs(W).sum(axis=0) > 0)
        idx = np.where(mem)[0]
        A, B = W[:-1, idx], W[1:, idx]
        A = (A - A.mean(0)) / (A.std(0) + 1e-12)
        B = (B - B.mean(0)) / (B.std(0) + 1e-12)
        L = A.T @ B / len(A)                                   # L[j, i]: j's week w-1 vs i's week w
        np.fill_diagonal(L, -np.inf)
        top = np.argsort(-L.T, axis=1)[:, :k]                  # for each i, its k best leaders
        out[me] = (idx, idx[top])
    return out


def tenk_links(p, role: str) -> pd.DataFrame:
    """Rows of (filer, target, filed) for the role ("customer", "supplier",
    "any"); "any" is made symmetric."""
    d = pd.read_csv(os.path.join("data", "research", "sec_10k_links.csv"), parse_dates=["filed"])
    if role != "any":
        d = d[d[f"n_{role}"] > 0]
    d = d[["filer", "target", "filed"]]
    if role == "any":
        d = pd.concat([d, d.rename(columns={"filer": "target", "target": "filer"})])
    return d


def tenk_link_sets(p, role: str, valid_days: int = 456) -> dict:
    """Month-end link sets from 10-Ks filed within the last 15 months (by
    filing date, so point-in-time)."""
    d = tenk_links(p, role)
    col = {t: i for i, t in enumerate(p.close.columns)}
    d = d[d["filer"].isin(col) & d["target"].isin(col) & (d["filer"] != d["target"])]
    out = {}
    for me in month_ends(p.dates):
        live = d[(d["filed"] <= me) & (d["filed"] > me - pd.Timedelta(days=valid_days))]
        if live.empty:
            continue
        mem = p.member.loc[me]
        g = {}
        for f, t in zip(live["filer"], live["target"]):
            if mem.get(f, False) and mem.get(t, False):
                g.setdefault(col[f], set()).add(col[t])
        if g:
            out[me] = {i: np.array(sorted(s)) for i, s in g.items()}
    return out


def link_signal(p, links: dict, horizon: int, kind: str) -> pd.DataFrame:
    """Mean linked-name log return over the last ``horizon`` days (kind="mom")
    or that minus the stock's own (kind="gap"), using each month's links for
    the days until the next month end."""
    lr = np.log1p(_ret(p).fillna(0.0))
    R = lr.rolling(horizon, min_periods=horizon).sum().to_numpy()
    out = np.full(R.shape, np.nan)
    mes = sorted(links)
    for a, b in zip(mes, mes[1:] + [p.dates[-1] + pd.Timedelta(days=1)]):
        rows = np.where((p.dates > a) & (p.dates <= b))[0]
        if not len(rows):
            continue
        L = links[a]
        if isinstance(L, tuple):
            idx, nb = L
            val = np.nanmean(R[np.ix_(rows, nb.ravel())].reshape(len(rows), *nb.shape), axis=2)
            out[np.ix_(rows, idx)] = val
        else:
            for i, nb in L.items():
                out[rows, i] = np.nanmean(R[np.ix_(rows, nb)], axis=1)
    S = pd.DataFrame(out, index=p.dates, columns=p.close.columns)
    if kind == "gap":
        S = S - pd.DataFrame(R, index=p.dates, columns=p.close.columns)
    return S.where(p.member)


# ---------------------------------------------------------------------------
# measures
# ---------------------------------------------------------------------------

def fama_macbeth(p, sig: pd.DataFrame, controls: dict[str, pd.DataFrame], horizon: int) -> dict:
    """Cross-sectional regressions every ``horizon`` days of the forward
    return rank on the signal rank and control ranks (all in [-0.5, 0.5])."""
    o = p.open.where(p.open.notna(), p.close).ffill()
    fwd = (o.shift(-(horizon + 1)) / o.shift(-1) - 1).where(p.member)     # next open to the open h days later
    rows = []
    for i in range(260, len(p.dates) - horizon, horizon):
        d = p.dates[i]
        y = fwd.iloc[i]
        cols = {"sig": sig.iloc[i], **{k: v.iloc[i] for k, v in controls.items()}}
        df = pd.DataFrame({"y": y, **cols}).dropna()
        if len(df) < 50:
            continue
        df = df.rank(pct=True) - 0.5
        X = np.column_stack([np.ones(len(df)), df.drop(columns="y").to_numpy()])
        b, *_ = np.linalg.lstsq(X, df["y"].to_numpy(), rcond=None)
        ic = df["y"].corr(df["sig"])
        rows.append({"date": d, "b_sig": b[1], "ic": ic, "n": len(df)})
    fm = pd.DataFrame(rows).set_index("date")

    def tstat(x):
        x = x.dropna()
        return float(x.mean() / x.std() * np.sqrt(len(x))) if len(x) > 2 else np.nan
    out = {}
    for name, sub in (("all", fm), ("design", fm.loc[:DESIGN_END]), ("holdout", fm.loc["2017":])):
        out[name] = {"b": float(sub["b_sig"].mean()), "t": tstat(sub["b_sig"]),
                     "ic": float(sub["ic"].mean()), "ic_t": tstat(sub["ic"]), "n_obs": float(sub["n"].mean()),
                     "periods": len(sub)}
    return out


def sleeve_stats(p, score: pd.DataFrame, top_n: int, every: int, bench) -> dict:
    from research.lab.alpha import alpha_stats
    from research.lab.us_sweep import run
    r = run(p, score.where(p.member), LONG, top_n=top_n, buffer=2.0, every=every)["_rets"]
    e = (1 + r).cumprod()
    out = {"mdd": float((e / e.cummax() - 1).min())}
    for name, x in (("all", r), ("design", r.loc[:DESIGN_END]), ("holdout", r.loc["2017":])):
        a = alpha_stats(x, bench)
        out[name] = {"alpha": a["alpha"], "t": a["t"], "cagr": a["cagr"], "bench": a["bench_cagr"], "beta": a["beta"]}
    return out, r


def resort(mom: pd.DataFrame, sig: pd.DataFrame, pool: int = 30) -> pd.DataFrame:
    """Keep the top ``pool`` by momentum, order them by the signal (the rest
    NaN). A name without links gets the neutral rank 0.5; momentum breaks ties."""
    rk = mom.rank(axis=1, ascending=False)
    s = sig.rank(axis=1, pct=True).fillna(0.5) + 1e-3 * mom.rank(axis=1, pct=True)
    return s.where(rk <= pool)


def fmt_sleeve(name: str, s: dict) -> str:
    d, h = s["design"], s["holdout"]
    return (f"{name:34s} alpha {s['all']['alpha']*100:5.1f}% t {s['all']['t']:5.2f} | 06-16 {d['alpha']*100:5.1f}% "
            f"t {d['t']:5.2f} cagr {d['cagr']*100:5.1f} | 17-26 {h['alpha']*100:5.1f}% t {h['t']:5.2f} "
            f"cagr {h['cagr']*100:5.1f} vs QQQ {h['bench']*100:5.1f} | mdd {s['mdd']*100:4.0f}%")


def part_signals(p, only: list[str] | None) -> dict:
    from research.lab.alpha import bench_returns
    from research.lab.us_sweep import _blend
    pr = factor_ranks(p, "us", ["overnight_mom", "rec_126", "mom_12_1"])
    mom = _blend(pr["overnight_mom"], pr["rec_126"]).where(p.member)
    c = p.close.ffill()
    own21 = (c / c.shift(21) - 1).where(p.member)
    own5 = (c / c.shift(5) - 1).where(p.member)
    bench = bench_returns(p.dates)
    cache = pickle.load(open(CACHE, "rb")) if os.path.exists(CACHE) else {}
    builders = {"corr_peer": lambda: corr_peer_links(p), "lead_peer": lambda: lead_peer_links(p)}
    if os.path.exists(os.path.join("data", "research", "sec_10k_links.csv")):
        for role in ("customer", "supplier", "any"):
            builders[f"10k_{role}"] = (lambda role=role: tenk_link_sets(p, role))
    res = {}
    base, base_r = sleeve_stats(p, mom, 10, 10, bench)
    print(fmt_sleeve("momentum top10 (reference)", base), flush=True)
    res["momentum_top10"] = base
    R = {"momentum_top10": base_r}
    for link, build in builders.items():
        if only and link not in only:
            continue
        if link not in cache:
            cache[link] = build()
            pickle.dump(cache, open(CACHE, "wb"))
        L = cache[link]
        cover = np.mean([len(v[0]) if isinstance(v, tuple) else len(v) for v in L.values()])
        hz = 5 if link == "lead_peer" else 21
        every = 5 if link == "lead_peer" else 21
        own = own5 if hz == 5 else own21
        print(f"\n== {link}: {len(L)} months, {cover:.0f} stocks with links on average", flush=True)
        for kind in ("mom", "gap"):
            sig = link_signal(p, L, hz, kind)
            fm = fama_macbeth(p, sig, {"own": own, "mom_12_1": pr["mom_12_1"], "overnight": pr["overnight_mom"]}, hz)
            print(f"  {link}_{kind}: FM b {fm['all']['b']:+.4f} t {fm['all']['t']:5.2f} | design t {fm['design']['t']:5.2f} "
                  f"| holdout t {fm['holdout']['t']:5.2f} | IC {fm['all']['ic']:+.4f} (t {fm['all']['ic_t']:.2f}) "
                  f"n {fm['all']['n_obs']:.0f}", flush=True)
            sl, r = sleeve_stats(p, sig, 20, every, bench)
            print("  " + fmt_sleeve(f"{link}_{kind} top20", sl), flush=True)
            rs, r2 = sleeve_stats(p, resort(mom, sig), 10, 10, bench)
            print("  " + fmt_sleeve(f"momentum top30 -> top10 by {kind}", rs), flush=True)
            res[f"{link}_{kind}"] = {"fm": fm, "sleeve": sl, "resort": rs,
                                     "corr_with_momentum": float(r.corr(base_r))}
            R[f"{link}_{kind}"] = r
            R[f"resort_{link}_{kind}"] = r2
    pd.DataFrame(R).to_pickle(os.path.join(RESULTS, "lead_lag_returns.pkl"))
    return res


# ---------------------------------------------------------------------------
# earnings spillover
# ---------------------------------------------------------------------------

def part_events(p) -> dict:
    """For each earnings day of member j, the abnormal return of its
    corr_peers (that did not report within 2 days) on the event day and over
    days 1-5, 1-10 and 1-21 after it, regressed on j's event-day abnormal return."""
    ev = pd.read_csv(os.path.join("data", "research", "sec_8k_earnings.csv"))
    ev["t"] = pd.to_datetime(ev["acceptanceDateTime"], utc=True).dt.tz_convert("America/New_York")
    # after 16:00 the market reacts on the next session
    ev["d"] = ev["t"].dt.normalize().dt.tz_localize(None) + pd.to_timedelta((ev["t"].dt.hour >= 16).astype(int), "D")
    dates = p.dates
    pos = dates.searchsorted(ev["d"])
    ev = ev[pos < len(dates)].assign(i=pos[pos < len(dates)])
    col = {t: k for k, t in enumerate(p.close.columns)}
    ev = ev[ev["ticker"].isin(col)]
    r = _ret(p)
    mkt = r.where(p.member).mean(axis=1)
    ab = r.sub(mkt, axis=0).to_numpy()
    cache = pickle.load(open(CACHE, "rb")) if os.path.exists(CACHE) else {}
    if "corr_peer" not in cache:
        cache["corr_peer"] = corr_peer_links(p)
        pickle.dump(cache, open(CACHE, "wb"))
    L = cache["corr_peer"]
    mes = np.array(sorted(L))
    rep = np.zeros(ab.shape, bool)
    for t, i in zip(ev["ticker"], ev["i"]):
        rep[max(i - 2, 0):i + 3, col[t]] = True
    rows = []
    for t, i in zip(ev["ticker"], ev["i"]):
        if i + 22 >= len(dates) or i < 300:
            continue
        me = mes[mes < dates[i]]
        if not len(me):
            continue
        idx, nb = L[me[-1]]
        j = col[t]
        w = np.where(idx == j)[0]
        if not len(w) or not p.member.iat[i, j]:
            continue
        a_j = ab[i, j]
        if not np.isfinite(a_j):
            continue
        for q in nb[w[0]]:
            if rep[i, q]:
                continue
            x = ab[i:i + 22, q]
            if not np.isfinite(x[0]):
                continue
            x = np.nan_to_num(x)
            rows.append((dates[i], t, p.close.columns[q], a_j, x[0], x[1:6].sum(), x[1:11].sum(), x[1:22].sum()))
    E = pd.DataFrame(rows, columns=["date", "announcer", "peer", "ar_j", "ar0", "car1_5", "car1_10", "car1_21"])
    E.to_pickle(os.path.join(RESULTS, "lead_lag_events.pkl"))
    out = {}
    print(f"\n{len(E)} (event, peer) pairs, {E['announcer'].nunique()} announcers")
    for name, sub in (("all", E), ("design", E[E["date"] <= DESIGN_END]), ("holdout", E[E["date"] > DESIGN_END])):
        res = {}
        for y in ("ar0", "car1_5", "car1_10", "car1_21"):
            # cluster by event date: average the slope within each month, t over months
            g = sub.groupby(sub["date"].dt.to_period("M"))
            sl = g.apply(lambda d: np.polyfit(d["ar_j"], d[y], 1)[0] if len(d) > 20 and d["ar_j"].std() > 0 else np.nan).dropna()
            res[y] = {"slope": float(sl.mean()), "t": float(sl.mean() / sl.std() * np.sqrt(len(sl)))}
        # the trade: peers of big positive (negative) surprises that barely moved themselves
        big = sub[sub["ar_j"].abs() > 0.05]
        lag = big[(np.sign(big["ar0"]) != np.sign(big["ar_j"])) | (big["ar0"].abs() < 0.01)]
        sgn = np.sign(lag["ar_j"])
        res["laggard_trade_car1_10"] = float((sgn * lag["car1_10"]).mean())
        res["laggard_trade_n"] = len(lag)
        res["laggard_trade_t"] = float((sgn * lag["car1_10"]).groupby(lag["date"].dt.to_period("M")).mean().pipe(
            lambda s: s.mean() / s.std() * np.sqrt(len(s))))
        out[name] = res
        print(f"  {name:8s} " + " | ".join(f"{y} slope {v['slope']:+.3f} t {v['t']:5.2f}" for y, v in res.items() if isinstance(v, dict))
              + f" | laggard trade 10d {res['laggard_trade_car1_10']*100:+.2f}% (t {res['laggard_trade_t']:.2f}, n {res['laggard_trade_n']})")
    return out


# ---------------------------------------------------------------------------
# AI chain case study (hindsight grouping, descriptive only)
# ---------------------------------------------------------------------------

def part_chain(p) -> dict:
    from research.lab.review_live_us import AI_CHAIN
    r = _ret(p)
    out = {}
    stages = {s: [t for t in ts if t in r.columns] for s, ts in AI_CHAIN.items()}
    W = pd.DataFrame({s: r[ts].mean(axis=1) for s, ts in stages.items()})
    mkt = r.where(p.member).mean(axis=1)
    W = W.sub(mkt, axis=0)                                   # stage return over the EW market
    wk = (1 + W).resample("W-FRI").prod() - 1
    for span in (("2016-01-01", "2022-12-31"), ("2023-01-01", "2026-09-24")):
        x = wk.loc[span[0]:span[1]]
        lead = pd.DataFrame(index=x.columns, columns=x.columns, dtype=float)
        for a in x.columns:
            for b in x.columns:
                lead.loc[a, b] = x[a].shift(1).corr(x[b])
        n = len(x)
        print(f"\nstage excess returns, weekly, {span[0][:4]}-{span[1][:4]} (n {n}, |corr| > {2/np.sqrt(n):.2f} is ~2 s.e.):"
              "\nrow = stage in week w-1, column = stage in week w")
        print(lead.round(2).to_string())
        out[f"{span[0][:4]}_{span[1][:4]}"] = lead.round(3).to_dict()
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--part", default="signals", choices=["signals", "events", "chain", "all"])
    ap.add_argument("--only", nargs="*")
    args = ap.parse_args()
    p, _ = load("us")
    out_path = os.path.join(RESULTS, "lead_lag.json")
    out = json.load(open(out_path)) if os.path.exists(out_path) else {}
    if args.part in ("signals", "all"):
        out.setdefault("signals", {}).update(part_signals(p, args.only))
    if args.part in ("events", "all"):
        out["events"] = part_events(p)
    if args.part in ("chain", "all"):
        out["chain"] = part_chain(p)
    json.dump(out, open(out_path, "w"), indent=1, default=float)


if __name__ == "__main__":
    main()