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

Writes data/research/us_fund_raw.pkl:
    {"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")
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)


def main() -> None:
    import sys
    with open(os.path.join("data", "research", "us_ohlcv.pkl"), "rb") as fh:
        tickers = [t for t in pickle.load(fh)["close"].columns]
    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 = 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()