"""
02_build_panel.py
Merge the semantic factor panel (annual_final) with realized stock performance.

Timing convention (avoids look-ahead bias):
  A 10-K for fiscal year Y is public by ~Feb-Mar of Y+1. We therefore form
  signals at the END OF JUNE of year Y+1 and measure the FORWARD 12-month
  return over July(Y+1) -> June(Y+2). This is the Fama-French portfolio
  convention: the signal is strictly known before the holding period starts.

Outputs:
  data/processed/monthly_returns.parquet   ticker, date, ret, mktcap_proxy
  data/processed/merged_panel.parquet      firm-year panel: AI scores (FY=Y) + fwd returns + controls
"""
import numpy as np
import pandas as pd

D = "data/processed"
px = pd.read_parquet(f"{D}/prices_monthly.parquet")
spy = pd.read_parquet(f"{D}/spy_monthly.parquet")
mc = pd.read_parquet(f"{D}/marketcap.parquet")
panel = pd.read_parquet("data/parquet/annual_final/annual_final.parquet")

# ---- monthly simple returns ----
px = px.sort_values(["ticker", "date"]).copy()
px["date"] = pd.to_datetime(px["date"]).dt.to_period("M").dt.to_timestamp("M")  # month-end stamp
px["ret"] = px.groupby("ticker")["adj_close"].pct_change()

spy["date"] = pd.to_datetime(spy["date"]).dt.to_period("M").dt.to_timestamp("M")
spy = spy.sort_values("date")
spy["spy_ret"] = spy["spy_adj_close"].pct_change()

# market-cap proxy = current shares * historical price (shares held constant approximation)
mc_map = mc.set_index("ticker")["shares"].to_dict()
px["shares"] = px["ticker"].map(mc_map)
px["mktcap_proxy"] = px["shares"] * px["adj_close"]
px.to_parquet(f"{D}/monthly_returns.parquet", index=False)

# ---- helper: cumulative return over [start, end] inclusive month-ends ----
def cum_ret(df_t, start, end):
    m = df_t[(df_t.date > start) & (df_t.date <= end)]
    if m.empty or m["ret"].isna().all():
        return np.nan
    r = (1 + m["ret"].fillna(0)).prod() - 1
    # require reasonable coverage
    return r if m["ret"].notna().sum() >= 6 else np.nan

spy_ix = spy.set_index("date")["spy_ret"]
def cum_spy(start, end):
    m = spy_ix[(spy_ix.index > start) & (spy_ix.index <= end)]
    return (1 + m.fillna(0)).prod() - 1 if len(m) else np.nan

rows = []
px_by = {t: d.sort_values("date") for t, d in px.groupby("ticker")}
for _, r in panel.iterrows():
    t = r["ticker"]; Y = int(r["fiscal_year"])
    d = px_by.get(t)
    if d is None:
        continue
    form = pd.Timestamp(year=Y+1, month=6, day=30) + pd.offsets.MonthEnd(0)
    end1 = pd.Timestamp(year=Y+2, month=6, day=30) + pd.offsets.MonthEnd(0)
    # controls measured AT formation (strictly known):
    mom_start = pd.Timestamp(year=Y, month=6, day=30) + pd.offsets.MonthEnd(0)   # 12m momentum to formation
    # forward 12m return
    fwd = cum_ret(d, form, end1)
    spy_fwd = cum_spy(form, end1)
    # momentum (past 12m, formation-12 .. formation, skip last month)
    mom = cum_ret(d, mom_start, form - pd.offsets.MonthEnd(1))
    # trailing vol & size at formation
    win = d[(d.date > mom_start) & (d.date <= form)]
    vol = win["ret"].std() * np.sqrt(12) if win["ret"].notna().sum() >= 6 else np.nan
    cap = d.loc[d.date <= form, "mktcap_proxy"].dropna()
    cap = cap.iloc[-1] if len(cap) else np.nan
    rows.append({"ticker": t, "fiscal_year": Y,
                 "fwd_ret": fwd, "spy_fwd": spy_fwd,
                 "fwd_ret_mktadj": fwd - spy_fwd if pd.notna(fwd) and pd.notna(spy_fwd) else np.nan,
                 "mom_12m": mom, "vol_12m": vol, "mktcap": cap,
                 "log_mktcap": np.log(cap) if pd.notna(cap) and cap > 0 else np.nan})

ret_panel = pd.DataFrame(rows)
merged = panel.merge(ret_panel, on=["ticker", "fiscal_year"], how="left")
merged.to_parquet(f"{D}/merged_panel.parquet", index=False)

n_ret = merged["fwd_ret"].notna().sum()
print(f"merged: {merged.shape}, firm-years with fwd return: {n_ret}")
print("fwd_ret by year (count):")
print(merged[merged.fwd_ret.notna()].groupby("fiscal_year").size().to_string())
print("\nfwd_ret describe:\n", merged["fwd_ret"].describe())