quant-trading/research/lab/fetch_us_fundamentals.py

"""
Fetch point-in-time US fundamentals from SEC EDGAR (XBRL companyfacts) for
every S&P 500 ever-member in the research price panel, plus Yahoo split
history (needed to turn split-adjusted prices back into market caps).

    uv run python -m research.lab.fetch_us_fundamentals
    uv run python -m research.lab.fetch_us_fundamentals --universe mid   # S&P 400 (IJH) ever-members
    uv run python -m research.lab.fetch_us_fundamentals --universe small # S&P 600 (IJR) ever-members

Writes data/research/us_fund_raw.pkl (us_{mid,small}_fund_raw.pkl for --universe):
    {"facts": DataFrame[ticker, tag, start, end, val, filed, form],
     "splits": DataFrame (dates x tickers, split ratio on split days)}

Every fact keeps its FIRST filing date: a value restated in a later filing is
not visible before that later date, so factors built from "filed" are
point-in-time. XBRL coverage starts in 2009-2011 depending on filer size.
"""

import os
import pickle
import time
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
import requests
import yfinance as yf

from research.lab.fetch_us import _proxy_env

OUT = os.path.join("data", "research", "us_fund_raw.pkl")
UNIVERSE_FUND = {"mid": "IJH", "small": "IJR"}       # --universe -> iShares fund
UA = {"User-Agent": "quant-research/0.1 gahow@example.com", "Accept-Encoding": "gzip"}
FORMS = {"10-K", "10-Q", "10-K/A", "10-Q/A", "10-KT", "20-F", "40-F"}

TAGS = [
    # income statement (durations)
    "Revenues", "RevenueFromContractWithCustomerExcludingAssessedTax",
    "RevenueFromContractWithCustomerIncludingAssessedTax", "SalesRevenueNet",
    "SalesRevenueGoodsNet", "SalesRevenueServicesNet", "RevenuesNetOfInterestExpense",
    "CostOfRevenue", "CostOfGoodsAndServicesSold", "CostOfGoodsSold", "CostOfServices",
    "GrossProfit", "OperatingIncomeLoss", "NetIncomeLoss", "ProfitLoss",
    "NetIncomeLossAvailableToCommonStockholdersBasic",
    "IncomeLossFromContinuingOperationsBeforeIncomeTaxesExtraordinaryItemsNoncontrollingInterest",
    "ResearchAndDevelopmentExpense", "SellingGeneralAndAdministrativeExpense",
    "InterestExpense", "IncomeTaxExpenseBenefit", "EarningsPerShareDiluted",
    "WeightedAverageNumberOfDilutedSharesOutstanding",
    "DepreciationDepletionAndAmortization", "DepreciationAndAmortization",
    # cash flow (year-to-date durations in 10-Qs)
    "NetCashProvidedByUsedInOperatingActivities",
    "NetCashProvidedByUsedInOperatingActivitiesContinuingOperations",
    "PaymentsToAcquirePropertyPlantAndEquipment", "PaymentsForRepurchaseOfCommonStock",
    "PaymentsOfDividends", "PaymentsOfDividendsCommonStock",
    # balance sheet (instants)
    "Assets", "AssetsCurrent", "Liabilities", "LiabilitiesCurrent",
    "LiabilitiesAndStockholdersEquity", "StockholdersEquity",
    "StockholdersEquityIncludingPortionAttributableToNoncontrollingInterest",
    "CashAndCashEquivalentsAtCarryingValue", "LongTermDebt", "LongTermDebtNoncurrent",
    "LongTermDebtCurrent", "DebtCurrent", "ShortTermBorrowings", "InventoryNet",
    "AccountsReceivableNetCurrent", "AccountsPayableCurrent", "PropertyPlantAndEquipmentNet",
    "Goodwill", "RetainedEarningsAccumulatedDeficit", "CommonStockSharesOutstanding",
]
DEI_TAGS = ["EntityCommonStockSharesOutstanding"]


def cik_map() -> dict[str, int]:
    r = requests.get("https://www.sec.gov/files/company_tickers.json", headers=UA, timeout=60)
    r.raise_for_status()
    return {v["ticker"].upper().replace(".", "-"): int(v["cik_str"]) for v in r.json().values()}


def _facts_for(ticker: str, cik: int) -> pd.DataFrame | None:
    url = f"https://data.sec.gov/api/xbrl/companyfacts/CIK{cik:010d}.json"
    for attempt in range(4):
        try:
            r = requests.get(url, headers=UA, timeout=120)
            if r.status_code == 404:
                return None
            r.raise_for_status()
            js = r.json()
            break
        except Exception as exc:  # noqa: BLE001 — retry transient SEC errors
            print(f"  {ticker}: {exc}; retry {attempt + 1}", flush=True)
            time.sleep(5 * (attempt + 1))
    else:
        return None
    rows = []
    for ns, tags in (("us-gaap", TAGS), ("dei", DEI_TAGS)):
        facts = js.get("facts", {}).get(ns, {})
        for tag in tags:
            if tag not in facts:
                continue
            for unit, xs in facts[tag]["units"].items():
                if unit not in ("USD", "shares", "USD/shares"):
                    continue
                for x in xs:
                    if x.get("form") not in FORMS:
                        continue
                    rows.append((ticker, tag, x.get("start"), x["end"], float(x["val"]),
                                 x["filed"], x["form"]))
    if not rows:
        return None
    df = pd.DataFrame(rows, columns=["ticker", "tag", "start", "end", "val", "filed", "form"])
    # first filing of each (tag, period) wins: later restatements are not
    # knowable on the original date
    df = df.sort_values("filed").drop_duplicates(["tag", "start", "end"], keep="first")
    return df


def fetch_splits(tickers: list[str], chunk: int = 60) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Split ratios (on split days, 0 elsewhere) and split-adjusted closes
    that are NOT dividend-adjusted; market cap = shares reported at c x
    close_sa[t] / (product of split ratios after c)."""
    splits, closes = [], []

    def run(batch):
        raw = yf.download(batch, start="2004-01-01", actions=True, auto_adjust=False,
                          progress=False, group_by="column", threads=True)
        lv = raw.columns.get_level_values(0)
        if "Stock Splits" in lv:
            splits.append(raw["Stock Splits"])
        if "Close" in lv:
            closes.append(raw["Close"].dropna(axis=1, how="all"))
        time.sleep(2)

    def cat(parts):
        s = pd.concat(parts, axis=1)
        return s.loc[:, ~s.columns.duplicated(keep="last")]

    for i in range(0, len(tickers), chunk):
        run(tickers[i:i + chunk])
    # Yahoo fails whole batches now and then; retry the missing names in small batches
    for attempt in range(3):
        have = set(cat(closes).columns)
        missing = [t for t in tickers if t not in have]
        if not missing:
            break
        print(f"  retrying {len(missing)} tickers", flush=True)
        for i in range(0, len(missing), 15):
            run(missing[i:i + 15])
    return cat(splits).fillna(0.0), cat(closes)


_SUFFIXES = {"INC", "CORP", "CORPORATION", "CO", "COMPANY", "LTD", "PLC", "HOLDINGS", "HOLDING",
             "GROUP", "THE", "CLASS", "CL", "A", "B", "NEW", "LLC", "LP", "SA", "NV", "TRUST"}


def _name_tokens(name: str) -> set[str]:
    words = "".join(ch if ch.isalnum() else " " for ch in str(name).upper()).split()
    return {w for w in words if w not in _SUFFIXES and len(w) > 1}


def mid_ciks(tickers: list[str], fund: str = "IJH") -> tuple[dict[str, int], pd.DataFrame]:
    """CIKs for S&P 400 (IJH) or S&P 600 (IJR) tickers. Candidates come from sec_sic.csv and
    mid_sec_sic.csv (earlier runs), then today's company_tickers.json; both
    map a ticker to whoever holds it NOW, so a candidate counts only if one of
    its SEC names (current or former) shares a word with an iShares holding
    name of that ticker. Returns the accepted map and a per-ticker audit."""
    import json
    from research.lab.fetch_us_mid import load_holdings
    base = os.path.join("data", "research")
    cand: dict[str, int] = {}
    for f in ("sec_sic.csv", "mid_sec_sic.csv"):
        d = pd.read_csv(os.path.join(base, f)).dropna(subset=["cik"])
        cand.update({t: int(c) for t, c in zip(d["ticker"], d["cik"]) if t not in cand})
    live = cik_map()
    for t in tickers:
        cand.setdefault(t, live.get(t, 0))
    names: dict[str, list[str]] = {}
    for f in ("sec_names.json", "mid_sec_names.json"):
        with open(os.path.join(base, f)) as fh:
            names.update(json.load(fh))
    for t in tickers:
        c = cand.get(t, 0)
        if c and str(c) not in names:
            r = requests.get(f"https://data.sec.gov/submissions/CIK{c:010d}.json", headers=UA, timeout=60)
            time.sleep(0.15)
            if r.ok:
                js = r.json()
                names[str(c)] = [js.get("name", "")] + [x.get("name", "") for x in js.get("formerNames", [])]
    hold = load_holdings()
    hnames = hold[hold["fund"] == fund].groupby("ticker")["name"].apply(lambda s: set(s.dropna()))
    rows, ok = [], {}
    for t in tickers:
        c = cand.get(t, 0)
        sec = names.get(str(c), []) if c else []
        hw = set().union(*[_name_tokens(n) for n in hnames.get(t, set())]) if t in hnames else set()
        sw = set().union(*[_name_tokens(n) for n in sec]) if sec else set()
        # token overlap, or the same leading letters once spacing and
        # punctuation are gone (MACYS / Macy's, SIRIUSXM / Sirius XM, VF / V F)
        squash = lambda n: "".join(w for w in "".join(ch if ch.isalnum() else " " for ch in str(n).upper()).split()
                                   if w not in _SUFFIXES)
        hs = {squash(n)[:5] for n in hnames.get(t, set())} - {""}
        ss = {squash(n)[:5] for n in sec} - {""}
        match = bool(hw & sw) or bool(hs & ss)
        if c and match:
            ok[t] = c
        rows.append({"ticker": t, "cik": c or None, "ishares_names": " | ".join(sorted(hnames.get(t, set()))),
                     "sec_names": " | ".join(sec), "accepted": bool(c and match)})
    return ok, pd.DataFrame(rows)


def main() -> None:
    import sys
    global OUT
    uni = sys.argv[sys.argv.index("--universe") + 1] if "--universe" in sys.argv else None
    if uni in UNIVERSE_FUND:
        from research.lab.panel_mid import load_us_mid
        OUT = os.path.join("data", "research", f"us_{uni}_fund_raw.pkl")
        tickers = list(load_us_mid(uni).close.columns)
        ciks, audit = mid_ciks(tickers, UNIVERSE_FUND[uni])
        audit.to_csv(os.path.join("data", "research", f"us_{uni}_fund_cik_audit.csv"), index=False)
        print(f"{len(ciks)}/{len(tickers)} {UNIVERSE_FUND[uni]} tickers have a name-checked CIK", flush=True)
    else:
        with open(os.path.join("data", "research", "us_ohlcv.pkl"), "rb") as fh:
            tickers = [t for t in pickle.load(fh)["close"].columns]
        ciks = None
    if "--yahoo-only" in sys.argv:        # refresh splits/closes of an existing file
        with open(OUT, "rb") as fh:
            saved = pickle.load(fh)
        _proxy_env()
        saved["splits"], saved["close_sa"] = fetch_splits(tickers)
        with open(OUT, "wb") as fh:
            pickle.dump(saved, fh)
        print(f"splits: {int((saved['splits'] > 0).sum().sum())} events; saved {OUT}")
        return
    ciks = ciks if ciks is not None else cik_map()
    todo = [(t, ciks[t]) for t in tickers if t in ciks]
    print(f"{len(todo)}/{len(tickers)} tickers map to a CIK", flush=True)

    frames, t0 = [], time.time()

    def job(tc):
        time.sleep(0.25)            # stays well under SEC's 10 requests/second
        return _facts_for(*tc)

    with ThreadPoolExecutor(max_workers=3) as ex:
        for k, df in enumerate(ex.map(job, todo), 1):
            if df is not None:
                frames.append(df)
            if k % 50 == 0:
                print(f"  {k}/{len(todo)} companies [{time.time() - t0:.0f}s]", flush=True)
    facts = pd.concat(frames, ignore_index=True)
    for c in ("start", "end", "filed"):
        facts[c] = pd.to_datetime(facts[c])
    print(f"{len(facts):,} facts for {facts['ticker'].nunique()} tickers", flush=True)

    _proxy_env()
    splits, close_sa = fetch_splits(tickers)
    print(f"splits: {int((splits > 0).sum().sum())} events", flush=True)
    with open(OUT, "wb") as fh:
        pickle.dump({"facts": facts, "splits": splits, "close_sa": close_sa,
                     "unmapped": sorted(set(tickers) - set(ciks))}, fh)
    print(f"Saved {OUT}")


if __name__ == "__main__":
    main()