fithia2/apps/tools/build_form4_pit_cache.py

"""Build leakage-safe daily Form 4 cluster parquet from SEC flat files.

The output is a daily same-day cluster cache keyed by filing date. Runtime code
uses this cache conservatively on the next trading day only.
"""
from __future__ import annotations

import argparse
import csv
import datetime as dt
import io
import math
import urllib.error
import urllib.request
import zipfile
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq

from apps.backtester.run import _build_merged_snapshot_store, load_manifest, resolve_config
from libs.common.logging import configure_logging

_DEFAULT_USER_AGENT = "fithia2-form4-cache/1.0 (local research; contact: dev@example.com)"


@dataclass(frozen=True)
class RawForm4Transaction:
    symbol: str
    filing_date: dt.date
    transaction_date: dt.date
    owner_cik: str
    owner_relationship: str
    owner_title: str
    shares: float
    price: float
    total_value: float
    shares_owned_following: float
    purchase_pct_of_holding: float


def _parse_date(value: str, *, is_end: bool = False) -> dt.date:
    parts = value.split("-")
    if len(parts) == 1 and len(value) == 4 and value.isdigit():
        year = int(value)
        return dt.date(year, 12, 31) if is_end else dt.date(year, 1, 1)
    if len(parts) == 2 and all(part.isdigit() for part in parts):
        year = int(parts[0])
        month = int(parts[1])
        if is_end:
            next_month = dt.date(year + (month // 12), (month % 12) + 1, 1)
            return next_month - dt.timedelta(days=1)
        return dt.date(year, month, 1)
    return dt.date.fromisoformat(value)


def _quarter_range(start_date: dt.date, end_date: dt.date) -> list[tuple[int, int]]:
    year = start_date.year
    quarter = (start_date.month - 1) // 3 + 1
    end_key = (end_date.year, (end_date.month - 1) // 3 + 1)
    quarters: list[tuple[int, int]] = []
    while (year, quarter) <= end_key:
        quarters.append((year, quarter))
        quarter += 1
        if quarter == 5:
            year += 1
            quarter = 1
    return quarters


def _quarter_zip_path(cache_dir: Path, year: int, quarter: int) -> Path:
    return cache_dir / f"{year}q{quarter}_form345.zip"


def _quarter_zip_url(year: int, quarter: int) -> str:
    return (
        "https://www.sec.gov/files/structureddata/data/"
        f"insider-transactions-data-sets/{year}q{quarter}_form345.zip"
    )


def _ensure_quarter_zip(cache_dir: Path, year: int, quarter: int, *, user_agent: str) -> Path | None:
    cache_dir.mkdir(parents=True, exist_ok=True)
    path = _quarter_zip_path(cache_dir, year, quarter)
    if path.exists() and path.stat().st_size > 0:
        return path
    request = urllib.request.Request(_quarter_zip_url(year, quarter), headers={"User-Agent": user_agent})
    try:
        with urllib.request.urlopen(request, timeout=90) as response:
            data = response.read()
    except urllib.error.HTTPError as exc:
        if exc.code == 404:
            return None
        raise
    path.write_bytes(data)
    return path


def _parse_sec_date(value: str | None) -> dt.date | None:
    if not value:
        return None
    try:
        return dt.datetime.strptime(value, "%d-%b-%Y").date()
    except ValueError:
        return None


def _coerce_float(value: Any) -> float | None:
    try:
        result = float(value)
    except (TypeError, ValueError):
        return None
    if math.isnan(result) or math.isinf(result):
        return None
    return result


def _is_officer_or_director(relationship: str, title: str) -> bool:
    combined = f"{relationship} {title}".strip().lower()
    return any(token in combined for token in ("director", "officer", "chief", "ceo", "cfo", "coo", "president", "chair"))


def _load_form4_transactions(
    *,
    cache_dir: Path,
    start_date: dt.date,
    end_date: dt.date,
    allowed_symbols: set[str],
    user_agent: str,
) -> tuple[list[RawForm4Transaction], list[str]]:
    transactions: list[RawForm4Transaction] = []
    skipped_quarters: list[str] = []

    for year, quarter in _quarter_range(start_date, end_date):
        path = _ensure_quarter_zip(cache_dir, year, quarter, user_agent=user_agent)
        if path is None:
            skipped_quarters.append(f"{year}Q{quarter}")
            continue
        with zipfile.ZipFile(path) as zf:
            submissions: dict[str, dict[str, Any]] = {}
            with zf.open("SUBMISSION.tsv") as handle:
                reader = csv.DictReader(io.TextIOWrapper(handle, encoding="utf-8", newline=""), delimiter="\t")
                for row in reader:
                    symbol = str(row.get("ISSUERTRADINGSYMBOL") or "").strip().upper()
                    if not symbol or symbol not in allowed_symbols:
                        continue
                    if str(row.get("DOCUMENT_TYPE") or "").strip().upper() != "4":
                        continue
                    filing_date = _parse_sec_date(row.get("FILING_DATE"))
                    if filing_date is None or filing_date < start_date or filing_date > end_date:
                        continue
                    submissions[str(row["ACCESSION_NUMBER"])] = {
                        "symbol": symbol,
                        "filing_date": filing_date,
                    }

            if not submissions:
                continue

            relationships: dict[str, list[tuple[str, str, str]]] = defaultdict(list)
            with zf.open("REPORTINGOWNER.tsv") as handle:
                reader = csv.DictReader(io.TextIOWrapper(handle, encoding="utf-8", newline=""), delimiter="\t")
                for row in reader:
                    accession = str(row["ACCESSION_NUMBER"])
                    if accession not in submissions:
                        continue
                    owner_cik = str(row.get("RPTOWNERCIK") or "").strip()
                    relationship = str(row.get("OFFICER_TITLE") or "").strip()
                    title = str(row.get("OTHER_TEXT") or "").strip()
                    relationships[accession].append((owner_cik, relationship, title))

            with zf.open("NONDERIV_TRANS.tsv") as handle:
                reader = csv.DictReader(io.TextIOWrapper(handle, encoding="utf-8", newline=""), delimiter="\t")
                for row in reader:
                    accession = str(row["ACCESSION_NUMBER"])
                    submission = submissions.get(accession)
                    if submission is None:
                        continue
                    transaction_code = str(row.get("TRANS_CODE") or "").strip().upper()
                    if transaction_code != "P":
                        continue
                    shares = _coerce_float(row.get("TRANS_SHARES"))
                    price = _coerce_float(row.get("TRANS_PRICEPERSHARE"))
                    if shares is None or price is None or shares <= 0 or price <= 0:
                        continue
                    transaction_date = _parse_sec_date(row.get("TRANS_DATE")) or submission["filing_date"]
                    shares_owned_following = _coerce_float(row.get("SHRS_OWND_FOLWNG_TRANS")) or 0.0
                    purchase_pct = 0.0
                    if shares_owned_following > 0:
                        purchase_pct = min(1.0, shares / shares_owned_following)
                    rels = relationships.get(accession) or [("", "", "")]
                    for owner_cik, relationship, title in rels:
                        transactions.append(
                            RawForm4Transaction(
                                symbol=submission["symbol"],
                                filing_date=submission["filing_date"],
                                transaction_date=transaction_date,
                                owner_cik=owner_cik,
                                owner_relationship=relationship,
                                owner_title=title,
                                shares=shares,
                                price=price,
                                total_value=shares * price,
                                shares_owned_following=shares_owned_following,
                                purchase_pct_of_holding=purchase_pct,
                            )
                        )

    return transactions, skipped_quarters


def _aggregate_daily_events(transactions: list[RawForm4Transaction]) -> list[dict[str, Any]]:
    grouped: dict[tuple[dt.date, str], list[RawForm4Transaction]] = defaultdict(list)
    for row in transactions:
        grouped[(row.filing_date, row.symbol)].append(row)

    events: list[dict[str, Any]] = []
    for (filing_date, symbol), rows in grouped.items():
        owner_keys = {row.owner_cik or f"{symbol}:{idx}" for idx, row in enumerate(rows)}
        purchase_values = [row.purchase_pct_of_holding for row in rows if row.purchase_pct_of_holding > 0]
        lag_values = [(row.filing_date - row.transaction_date).days for row in rows]
        events.append(
            {
                "symbol": symbol,
                "filing_date": filing_date.isoformat(),
                "as_of_date": filing_date.isoformat(),
                "total_value": round(sum(row.total_value for row in rows), 2),
                "owner_count": len(owner_keys),
                "transaction_count": len(rows),
                "event_day_count": 1,
                "max_purchase_pct": max(purchase_values) if purchase_values else 0.0,
                "median_purchase_pct": sorted(purchase_values)[len(purchase_values) // 2] if purchase_values else 0.0,
                "weighted_purchase_pct": max(purchase_values) if purchase_values else 0.0,
                "max_lag_days": max(lag_values) if lag_values else None,
                "min_lag_days": min(lag_values) if lag_values else None,
                "has_officer_or_director": any(
                    _is_officer_or_director(row.owner_relationship, row.owner_title)
                    for row in rows
                ),
            }
        )

    events.sort(key=lambda row: (row["filing_date"], row["symbol"]))
    return events


def _allowed_symbols_from_manifest(manifest_path: str, start_date: dt.date, end_date: dt.date) -> set[str]:
    manifest = load_manifest(manifest_path)
    config = resolve_config(manifest)
    store = _build_merged_snapshot_store(
        manifest,
        config,
        snapshot_dir_override=None,
    ).slice_by_date_range(start_date, end_date)
    return {str(symbol).upper() for symbol in store._bars.keys()}


def main() -> None:
    parser = argparse.ArgumentParser(description="Build daily Form 4 same-day cluster parquet")
    parser.add_argument("--config", required=True, help="Experiment manifest JSON path")
    parser.add_argument("--start", required=True, help="Start date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument("--end", required=True, help="End date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument("--cache-dir", default="data/cache/form4", help="SEC zip cache dir")
    parser.add_argument("--output", default="data/reference/form4_daily_events_pit.parquet", help="Output parquet path")
    parser.add_argument("--user-agent", default=_DEFAULT_USER_AGENT, help="SEC User-Agent header")
    args = parser.parse_args()

    configure_logging("INFO")
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    allowed_symbols = _allowed_symbols_from_manifest(args.config, start_date, end_date)
    transactions, skipped_quarters = _load_form4_transactions(
        cache_dir=Path(args.cache_dir),
        start_date=start_date,
        end_date=end_date,
        allowed_symbols=allowed_symbols,
        user_agent=args.user_agent,
    )
    events = _aggregate_daily_events(transactions)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    table = pa.Table.from_pylist(events)
    pq.write_table(table, str(output_path))
    print(
        {
            "output": str(output_path),
            "symbols": len({row["symbol"] for row in events}),
            "events": len(events),
            "transactions": len(transactions),
            "skipped_quarters": skipped_quarters,
        }
    )


if __name__ == "__main__":
    main()