"""
Point-in-time fundamental factors for the US panel from SEC EDGAR facts
(research/lab/fetch_us_fundamentals.py).

A value is usable from the trading day AFTER its first filing date, so a
factor on date t only uses statements filed before t. Durations become
quarterly values (direct 3-month facts, or differences of year-to-date facts
that share a fiscal-year start) and trailing-twelve-month sums; instants are
taken as of their period end.

    from research.lab import fundamentals as fu
    F = fu.load(panel)          # {name: dates x tickers}, cached
"""

from __future__ import annotations

import os
import pickle
import time

import numpy as np
import pandas as pd

RAW = os.path.join("data", "research", "us_fund_raw.pkl")
CACHE = os.path.join("data", "research", "us_fund_factors.pkl")

DURATION = {
    "revenue": ["Revenues", "RevenueFromContractWithCustomerExcludingAssessedTax",
                "RevenueFromContractWithCustomerIncludingAssessedTax", "SalesRevenueNet",
                "RevenuesNetOfInterestExpense", "SalesRevenueGoodsNet", "SalesRevenueServicesNet"],
    "cogs": ["CostOfRevenue", "CostOfGoodsAndServicesSold", "CostOfGoodsSold", "CostOfServices"],
    "gross_profit": ["GrossProfit"],
    "op_income": ["OperatingIncomeLoss"],
    "net_income": ["NetIncomeLoss", "ProfitLoss", "NetIncomeLossAvailableToCommonStockholdersBasic"],
    "cfo": ["NetCashProvidedByUsedInOperatingActivities",
            "NetCashProvidedByUsedInOperatingActivitiesContinuingOperations"],
    "capex": ["PaymentsToAcquirePropertyPlantAndEquipment"],
    "rnd": ["ResearchAndDevelopmentExpense"],
    "buyback": ["PaymentsForRepurchaseOfCommonStock"],
    "dividends": ["PaymentsOfDividendsCommonStock", "PaymentsOfDividends"],
}
INSTANT = {
    "assets": ["Assets"],
    "assets_cur": ["AssetsCurrent"],
    "liab": ["Liabilities"],
    "liab_cur": ["LiabilitiesCurrent"],
    "equity": ["StockholdersEquity",
               "StockholdersEquityIncludingPortionAttributableToNoncontrollingInterest"],
    "ltd": ["LongTermDebtNoncurrent", "LongTermDebt"],
    "shares": ["EntityCommonStockSharesOutstanding", "CommonStockSharesOutstanding"],
}


# ---------------------------------------------------------------------------
# fact tables → per-ticker event series
# ---------------------------------------------------------------------------

def _pick(facts: pd.DataFrame, tags: list[str]) -> pd.DataFrame:
    """One value per (ticker, start, end): the highest-priority tag present."""
    sub = facts[facts["tag"].isin(tags)].copy()
    sub["prio"] = sub["tag"].map({t: i for i, t in enumerate(tags)})
    sub = sub.sort_values(["prio", "filed"])
    return sub.drop_duplicates(["ticker", "start", "end"], keep="first").drop(columns="prio")


def _quarters(df: pd.DataFrame) -> pd.DataFrame:
    """Quarterly values from duration facts of one ticker: direct ~3-month
    facts, else differences of year-to-date facts sharing a start date."""
    df = df.dropna(subset=["start"])
    days = (df["end"] - df["start"]).dt.days
    direct = df[(days >= 75) & (days <= 105)][["end", "val", "filed"]]
    derived = []
    for _, g in df[days > 105].groupby("start"):
        g = g.sort_values("end")
        # pair each YTD fact with the one ending ~3 months earlier (same start)
        allg = pd.concat([g, df[(df["start"] == g["start"].iloc[0]) & (days.loc[df.index] <= 105)]])
        allg = allg.sort_values("end").drop_duplicates("end")
        e = allg["end"].values
        v = allg["val"].values
        f = allg["filed"].values
        for k in range(1, len(allg)):
            gap = (e[k] - e[k - 1]).astype("timedelta64[D]").astype(int)
            if 75 <= gap <= 105:
                derived.append((e[k], v[k] - v[k - 1], max(f[k], f[k - 1])))
    der = pd.DataFrame(derived, columns=["end", "val", "filed"])
    q = pd.concat([direct.assign(src=0), der.assign(src=1)]).sort_values(["end", "src", "filed"])
    q = q.drop_duplicates("end", keep="first").drop(columns="src").sort_values("end")
    # the first quarter of a fiscal year never has a same-start predecessor;
    # its direct fact covers it
    return q.reset_index(drop=True)


def _ttm(q: pd.DataFrame) -> pd.DataFrame:
    """Sum of four consecutive quarters (ends 75-105 days apart)."""
    if len(q) < 4:
        return q.iloc[:0]
    e = q["end"].values
    gaps = np.diff(e).astype("timedelta64[D]").astype(int)
    ok_gap = np.r_[False, (gaps >= 75) & (gaps <= 105)]
    run = pd.Series(ok_gap).rolling(3).sum().values == 3       # 3 good gaps → 4 quarters
    vals = q["val"].rolling(4).sum().values
    filed = pd.to_datetime(q["filed"].astype("int64").rolling(4).max().values)
    out = pd.DataFrame({"end": e, "val": vals, "filed": filed})
    return out[run].reset_index(drop=True)


def _asof_panel(events: dict[str, pd.DataFrame], dates: pd.DatetimeIndex,
                tickers: list[str], lag: int = 0) -> pd.DataFrame:
    """Daily panel of the value with the latest period end among facts filed
    strictly before each date (plus ``lag`` extra periods back: lag=4 gives
    the value four quarters earlier than the latest)."""
    out = pd.DataFrame(np.nan, index=dates, columns=tickers)
    for t, ev in events.items():
        if t not in out.columns or ev.empty:
            continue
        ev = ev.sort_values("filed")
        # as of each filing: the known history sorted by end
        rows = []
        known: dict = {}
        for end, val, filed in zip(ev["end"].values, ev["val"].values, ev["filed"].values):
            known[end] = val
            ends = sorted(known)
            if len(ends) > lag:
                rows.append((filed, known[ends[-1 - lag]]))
        if not rows:
            continue
        s = pd.DataFrame(rows, columns=["filed", "val"]).drop_duplicates("filed", keep="last")
        s = s.set_index("filed")["val"]
        # usable from the next session after the filing date
        pos = dates.searchsorted(s.index, side="right")
        ser = pd.Series(np.nan, index=dates)
        valid = pos < len(dates)
        ser.iloc[pos[valid]] = s.values[valid]
        out[t] = ser.ffill().values
    return out


def _series_by_ticker(df: pd.DataFrame) -> dict[str, pd.DataFrame]:
    return {t: g for t, g in df.groupby("ticker")}


# ---------------------------------------------------------------------------
# build
# ---------------------------------------------------------------------------

def build(p) -> dict[str, pd.DataFrame]:
    with open(RAW, "rb") as fh:
        raw = pickle.load(fh)
    facts = raw["facts"]
    dates, tickers = p.dates, list(p.close.columns)
    t0 = time.time()
    Q, TTM, INS = {}, {}, {}
    for c, tags in DURATION.items():
        picked = _pick(facts, tags)
        q_ev, ttm_ev = {}, {}
        for t, g in picked.groupby("ticker"):
            q = _quarters(g)
            q_ev[t] = q
            ttm = _ttm(q)
            # fiscal-year facts give TTM at year ends directly
            days = (g["end"] - g["start"]).dt.days
            fy = g[(days >= 350) & (days <= 380)][["end", "val", "filed"]]
            ttm_ev[t] = (pd.concat([fy.assign(src=0), ttm.assign(src=1)])
                         .sort_values(["end", "src"]).drop_duplicates("end").drop(columns="src"))
        Q[c], TTM[c] = q_ev, ttm_ev
        print(f"  {c}: {len(q_ev)} tickers [{time.time() - t0:.0f}s]", flush=True)
    for c, tags in INSTANT.items():
        if c == "shares":
            # cover-page and balance-sheet share counts have different dates;
            # mixing them would break "four reports back", so use one tag per ticker
            sub = facts[facts["tag"].isin(tags)]
            INS[c] = {}
            for t, g in sub.groupby("ticker"):
                dei = g[g["tag"] == tags[0]]
                use = dei if len(dei) >= 8 else g[g["tag"] == tags[1]]
                INS[c][t] = use.drop_duplicates("end")[["end", "val", "filed"]]
            continue
        picked = _pick(facts, tags)
        INS[c] = {t: g[["end", "val", "filed"]] for t, g in picked.groupby("ticker")}

    def ttm(c, lag=0):
        return _asof_panel(TTM[c], dates, tickers, lag)

    def qv(c, lag=0):
        return _asof_panel(Q[c], dates, tickers, lag)

    def ins(c, lag=0):
        return _asof_panel(INS[c], dates, tickers, lag)

    F: dict[str, pd.DataFrame] = {}
    rev, ni, cfo = ttm("revenue"), ttm("net_income"), ttm("cfo")
    gp = ttm("gross_profit")
    cogs = ttm("cogs")
    gp = gp.combine_first(rev - cogs)
    opi = ttm("op_income")
    capex = ttm("capex")
    assets, equity, liab = ins("assets"), ins("equity"), ins("liab")
    assets_4 = ins("assets", 4)
    print(f"  panels [{time.time() - t0:.0f}s]", flush=True)
    pos = lambda x: x.where(x > 0)   # noqa: E731

    F["gp_a"] = gp / pos(assets)                                   # Novy-Marx 2013
    F["op_a"] = opi / pos(assets)
    F["cfo_a"] = cfo / pos(assets)                                 # cash profitability
    F["roe"] = ni / pos(equity)
    F["roa"] = ni / pos(assets)
    F["accruals"] = -(ni - cfo) / pos(assets)                      # Sloan 1996 (low is good)
    F["asset_growth"] = -(assets / pos(assets_4) - 1)              # Cooper-Gulen-Schill 2008
    F["leverage"] = -(liab / pos(assets))
    F["gross_margin"] = gp / pos(rev)
    rev_4 = ttm("revenue", 4)
    F["rev_growth"] = rev / pos(rev_4) - 1
    gm_4 = ttm("gross_profit", 4).combine_first(rev_4 - ttm("cogs", 4)) / pos(rev_4)
    F["gm_change"] = F["gross_margin"] - gm_4
    ni_4 = ttm("net_income", 4)
    F["roa_change"] = F["roa"] - ni_4 / pos(assets_4)

    # SUE: latest quarterly change vs same quarter last year, scaled by the
    # std of the last 8 such changes (seasonal random walk, Bernard-Thomas)
    for c, name in (("net_income", "sue"), ("revenue", "sue_rev")):
        chg = {}
        for t, q in Q[c].items():
            q = q.sort_values("end").reset_index(drop=True)
            if len(q) < 8:
                continue
            e = q["end"].values
            v = q["val"].values
            d = np.full(len(q), np.nan)
            for k in range(4, len(q)):
                gap = (e[k] - e[k - 4]).astype("timedelta64[D]").astype(int)
                if 350 <= gap <= 380:
                    d[k] = v[k] - v[k - 4]
            ds = pd.Series(d)
            sd = ds.rolling(8, min_periods=4).std()
            sue = (ds / sd.replace(0, np.nan)).values
            chg[t] = pd.DataFrame({"end": e, "val": sue, "filed": q["filed"].values}).dropna()
        F[name] = _asof_panel(chg, dates, tickers).clip(-10, 10)

    # market cap. Yahoo's split-adjusted close is raw / F(t), where F(t) is the
    # product of split ratios after t, and shares reported at c grow by the
    # splits in (c, t], so cap_t = shares_c * F(c) * close_sa_t.
    splits = raw["splits"].reindex(columns=tickers).fillna(0.0)
    close_sa = raw["close_sa"].reindex(index=dates, columns=tickers).ffill()
    ratio = splits.replace(0.0, 1.0)
    # F(c) = product of split ratios strictly after c
    after = ratio.iloc[::-1].cumprod().iloc[::-1].shift(-1).fillna(1.0)
    after = after.reindex(dates, method="ffill").fillna(1.0)
    sh_ev = {}
    for t, g in INS["shares"].items():
        if t not in after.columns:
            continue
        ff = after[t].reindex(g["end"].values, method="ffill").fillna(1.0).values
        sh_ev[t] = g.assign(val=g["val"].values * ff)          # shares on today's split basis
    shares_now = _asof_panel(sh_ev, dates, tickers)
    mcap = shares_now * close_sa
    mcap = mcap.where(mcap > 1e8)
    F["mcap"] = mcap
    F["ep"] = ni / mcap
    F["cfp"] = cfo / mcap
    F["sp"] = rev / mcap
    F["bm"] = equity / mcap
    F["fcfp"] = (cfo - capex.fillna(0)) / mcap
    payout = ttm("buyback").fillna(0) + ttm("dividends").fillna(0)
    F["shareholder_yield"] = payout / mcap
    sh_4 = _asof_panel(sh_ev, dates, tickers, lag=4)
    F["net_issuance"] = -np.log(shares_now / sh_4).clip(-1, 1)

    # Piotroski F-score (9 binary signals)
    ltd, ac, lc = ins("ltd"), ins("assets_cur"), ins("liab_cur")
    ltd_4, ac_4, lc_4 = ins("ltd", 4), ins("assets_cur", 4), ins("liab_cur", 4)
    sig = [
        F["roa"] > 0, F["cfo_a"] > 0, F["roa_change"] > 0, cfo > ni,
        (ltd / pos(assets)) < (ltd_4 / pos(assets_4)),
        (ac / pos(lc)) > (ac_4 / pos(lc_4)),
        F["net_issuance"] >= -0.005,
        F["gm_change"] > 0,
        (rev / pos(assets)) > (rev_4 / pos(assets_4)),
    ]
    avail = F["roa"].notna() & F["cfo_a"].notna()
    F["fscore"] = sum(s.astype(float) for s in sig).where(avail)
    print(f"  factors built [{time.time() - t0:.0f}s]", flush=True)
    return {k: v.astype("float32") for k, v in F.items()}


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)
    F = build(p)
    with open(CACHE, "wb") as fh:
        pickle.dump(F, fh, protocol=4)
    return F