"""
Earnings-announcement return (EAR) factor: the market's reaction to the
latest earnings release, which keeps drifting afterwards (Chan, Jegadeesh &
Lakonishok 1996; Brandt, Kishore, Santa-Clara & Venkatachalam 2008 find the
drift after the announcement *return* stronger and longer-lived than after
the earnings surprise).

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

Announcement dates: 8-K filings with Item 2.02 (Results of Operations),
data/research/sec_8k_earnings.csv (written by research/lab/fetch_lazy_prices.py
from the SEC submissions JSON). A release after the close on day d is
reacted to on d+1, so the event window is the three sessions d-1..d+1:

    ear = sum of (stock return - equal-weight member return) over d-1, d, d+1

It is known at the close of d+1, the score is dated then (the harness trades
at the next open) and held until the next release, at most 100 trading days.
Also: ``ear_z`` = ear divided by the stock's 60-day residual vol (a surprise in
units of normal noise), and ``ear_count`` = the sum of the last four signs.
"""

from __future__ import annotations

import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.factors import _mkt, _ret

warnings.filterwarnings("ignore")
SRC = os.path.join("data", "research", "sec_8k_earnings.csv")
CACHE = os.path.join("data", "research", "earnings_factors.pkl")
MAX_AGE = 100


def announcement_days(p) -> pd.DataFrame:
    """Bool panel: an Item 2.02 8-K was filed that day (mapped to the next
    trading day when filed on a holiday or weekend)."""
    ev = pd.read_csv(SRC, parse_dates=["filingDate"])
    ev = ev[ev["ticker"].isin(p.close.columns)]
    pos = p.dates.searchsorted(ev["filingDate"].values)
    ev = ev.assign(pos=pos)[pos < len(p.dates)]
    A = np.zeros(p.close.shape, bool)
    col = {c: i for i, c in enumerate(p.close.columns)}
    A[ev["pos"].to_numpy(), ev["ticker"].map(col).to_numpy()] = True
    return pd.DataFrame(A, index=p.dates, columns=p.close.columns)


def _hold(values: pd.DataFrame, known: pd.DataFrame, max_age: int) -> pd.DataFrame:
    """Carry each event's value forward from the day it is known, for at most
    ``max_age`` trading days."""
    idx = np.broadcast_to(np.arange(len(values))[:, None], values.shape)
    last_pos = pd.DataFrame(np.where(known.values, idx, np.nan), index=values.index,
                            columns=values.columns).ffill()
    return values.where(known).ffill().mask(idx - last_pos.values > max_age)


def build(p) -> dict[str, pd.DataFrame]:
    A = announcement_days(p)
    ab = _ret(p).sub(_mkt(p), axis=0)                          # abnormal daily return
    car = ab.shift(1) + ab + ab.shift(-1)                      # d-1..d+1, indexed at d
    known = A.shift(1, fill_value=False)                       # known at the close of d+1
    vol = ab.rolling(60, min_periods=40).std().shift(2)        # pre-event noise
    ear = _hold(car.shift(1), known, MAX_AGE)
    ear_z = _hold((car / vol).shift(1), known, MAX_AGE)
    sign = np.sign(car.shift(1)).where(known)
    count = pd.DataFrame(np.nan, index=p.dates, columns=p.close.columns)
    for c in count.columns:
        x = sign[c].dropna()
        count.loc[x.index, c] = x.rolling(4, min_periods=2).sum()
    ear_count = _hold(count, known & count.notna(), MAX_AGE)
    return {"ear": ear, "ear_z": ear_z, "ear_count": ear_count, "announce": A}


def load(p, rebuild: bool = False) -> dict[str, pd.DataFrame]:
    if os.path.exists(CACHE) and not rebuild:
        with open(CACHE, "rb") as fh:
            return pickle.load(fh)
    out = build(p)
    with open(CACHE, "wb") as fh:
        pickle.dump(out, fh, protocol=4)
    return out


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

    p, _ = load_panel("us")
    E = load(p, rebuild=True)
    A = E["announce"] & p.member
    print("announcements among members per year:", A.sum(axis=1).groupby(p.dates.year).sum().to_dict())
    print("members with a live ear value (median per day by year):",
          E["ear"].where(p.member).notna().sum(axis=1).groupby(p.dates.year).median().to_dict())
    fwd = forward_returns(p, 21).where(p.member)
    bench = bench_returns(p.dates)
    pr = factor_ranks(p, "us")
    mom = _blend(pr["overnight_mom"], pr["rec_126"])
    for name in ("ear", "ear_z", "ear_count"):
        rk = E[name].where(p.member).rank(axis=1, pct=True)
        ic = rk.corrwith(fwd.rank(axis=1, pct=True), axis=1).iloc[::21].dropna()
        for a, b in (("2006", "2016"), ("2017", "2026")):
            x = ic.loc[a:b]
            print(f"{name} IC21 {a}-{b}: {x.mean():+.4f} t {x.mean() / x.std() * np.sqrt(len(x)):+.2f}")
        for label, sc, kw in ((f"{name} top20 e21", rk, dict(top_n=20, every=21)),
                              (f"{name} top20 e10", rk, dict(top_n=20, every=10)),
                              (f"momentum + {name} top10 e10", _blend(mom, rk), dict(top_n=10, every=10))):
            st = run(p, sc.where(p.member), [("2006-09-27", "2026-09-24")], buffer=2.0, **kw)
            r = st["_rets"]
            print(fmt_alpha(label, alpha_stats(r, bench)), "| holdout",
                  fmt_alpha("", alpha_stats(r.loc["2017":], bench)).split("alpha", 1)[1][:40], flush=True)
    st = run(p, mom.where(p.member), [("2006-09-27", "2026-09-24")], top_n=10, buffer=2.0, every=10)
    print(fmt_alpha("momentum alone top10 e10", alpha_stats(st["_rets"], bench)))


if __name__ == "__main__":
    main()