"""Institutional-flow factor scan from the DERA 13F aggregates.

Per quarter, total institutional shares per CUSIP are aggregated by
fetch_13f.py. Here they map to tickers (issuer name -> SEC title map) and
become the quarterly factor:

  inst_flow  = shares(q) / shares(q-1) - 1        (net institutional accumulation)
  filer_chg  = n_filers(q) - n_filers(q-1)        (breadth of accumulation)

PIT: the 13F deadline is 45 days after quarter end; the factor is applied
conservatively from quarter_end + 60 calendar days.

    PYTHONPATH=. uv run python -m research.lab.inst_flow [us|midsmall]
"""
from __future__ import annotations

import glob
import json
import os
import pickle
import re
import sys

import numpy as np
import pandas as pd

DIR = os.path.join("data", "research", "dera13f")
TICKERS_CACHE = "data/research/company_tickers.json"
PUBLISH_LAG_DAYS = 60

_SUFFIX = re.compile(r"[,\s]*(inc|incorporated|corp|corporation|co|company|ltd|limited|plc|"
                     r"holdings?|group|the|nv|sa|ag|llc|lp|trust|reit|new|cl|a|b)[.\s]*$", re.I)


def norm_name(name: str) -> str:
    if not isinstance(name, str):
        return ""
    n = name.upper().strip().rstrip(".")
    n = re.sub(r"[^A-Z0-9 ]+", " ", n)
    n = re.sub(r"\s+", " ", n)
    prev = None
    while prev != n:
        prev = n
        n = _SUFFIX.sub("", n).strip()
    return n


def load_quarter(path: str) -> pd.DataFrame:
    d = pd.read_csv(path, dtype={"CUSIP": str})
    d = d[d["CUSIP"].notna() & (d["CUSIP"] != "000000000")]
    d["issuer"] = d["issuer"].astype(str)
    d["norm"] = d["issuer"].map(norm_name)
    return d


def main() -> None:
    which = sys.argv[1] if len(sys.argv) > 1 else "us"
    files = sorted(glob.glob(os.path.join(DIR, "agg_*.csv")))
    if not files:
        raise SystemExit("no agg files; run fetch_13f first")
    print(f"quarters: {len(files)}")
    raw = json.load(open(TICKERS_CACHE))
    name2t = {}
    for v in raw.values():
        name2t.setdefault(norm_name(str(v["title"])), str(v["ticker"]).upper().replace(".", "-"))
    # combine quarters: cusip -> norm name from latest, and shares series
    shares = {}
    filers = {}
    name_by_cusip = {}
    for path in files:
        q = os.path.basename(path)[4:-4]  # YYYYqQ
        d = load_quarter(path)
        d = d[d["shares"].fillna(0) > 0]
        shares[q] = d.set_index("CUSIP")["shares"]
        filers[q] = d.set_index("CUSIP")["n_filers"]
        for c, n in d.set_index("CUSIP")["norm"].items():
            name_by_cusip[c] = n
    qs = [os.path.basename(p)[4:-4] for p in files]
    matched, ticker_by_cusip = 0, {}
    for c, n in name_by_cusip.items():
        t = name2t.get(n)
        if t:
            ticker_by_cusip[c] = t
            matched += 1
    print(f"cusip->ticker matched: {matched}/{len(name_by_cusip)}")

    # factor frame on quarter end dates (apply from +PUBLISH_LAG_DAYS)
    rows = {}
    prev = None
    for q in qs:
        if prev is None:
            prev = q
            continue
        s1 = shares[prev].reindex(shares[q].index)
        flow = shares[q] / s1.replace(0.0, np.nan) - 1.0
        d_fil = filers[q] - filers[prev].reindex(filers[q].index)
        # quarter-end date
        y, qq = int(q[:4]), int(q[-1])
        qend = pd.Timestamp(y, qq * 3, 1) + pd.offsets.MonthEnd(0)
        avail = qend + pd.Timedelta(days=PUBLISH_LAG_DAYS)
        t = pd.Series({ticker_by_cusip[c]: f for c, f in flow.items() if c in ticker_by_cusip and np.isfinite(f)})
        tc = pd.Series({ticker_by_cusip[c]: f for c, f in d_fil.items() if c in ticker_by_cusip and np.isfinite(f)})
        if len(t):
            rows[avail] = (t, tc)
        prev = q
    print(f"quarters with factors: {len(rows)}")

    if which == "us":
        with open("data/research/us.panel.pkl", "rb") as fh:
            p = pickle.load(fh)
    else:
        from research.lab.panel_mid import load_us_mid
        p = load_us_mid(which)
    close, member = p.close, p.member.reindex_like(p.close).fillna(False)
    fwd21 = close.shift(-21) / close - 1.0
    # broadcast each quarterly factor until the next availability date
    fac = pd.DataFrame(np.nan, index=close.index, columns=close.columns)
    facf = fac.copy()
    dates_avail = sorted(rows)
    for i, dt in enumerate(dates_avail):
        t, tc = rows[dt]
        end = dates_avail[i + 1] if i + 1 < len(dates_avail) else close.index[-1]
        sel = (close.index >= dt) & (close.index < end)
        for name, ser in (("flow", t), ("filers", tc)):
            frame = fac if name == "flow" else facf
            frame.loc[sel, frame.columns.isin(ser.index)] = ser.reindex(
                frame.columns[frame.columns.isin(ser.index)]).values
    months = fwd21.index.to_period("M")
    for name, frame in (("inst_flow", fac), ("filer_chg", facf)):
        f = frame.where(member)
        ics = []
        for per in sorted(set(months)):
            idx = months == per
            xdf = f.loc[idx]
            if xdf.notna().sum(axis=1).max() < 50:
                continue
            ydf = fwd21.loc[idx]
            mask = xdf.notna().to_numpy() & np.isfinite(ydf.reindex(columns=xdf.columns).to_numpy())
            if mask.sum() < 100:
                continue
            xr = xdf.rank(axis=1).to_numpy()[mask]
            yr = ydf.reindex(columns=xdf.columns).where(mask).rank(axis=1).to_numpy()[mask]
            xm, ym = xr - xr.mean(), yr - yr.mean()
            denom = np.sqrt((xm ** 2).sum() * (ym ** 2).sum())
            if denom > 0:
                ics.append(float((xm * ym).sum() / denom))
        ic = np.array(ics)
        print(f"{name:10s} IC21 {ic.mean():+.4f}  t={ic.mean()/ic.std()*np.sqrt(len(ic)):+.2f}  months={len(ic)}")


if __name__ == "__main__":
    main()