fithia2/libs/backtest/form4_calendar.py

"""Point-in-time Form 4 cluster helpers for leakage-safe residual cash sleeves."""
from __future__ import annotations

import datetime as dt
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Iterable

import pyarrow.parquet as pq

from libs.common.logging import get_logger

logger = get_logger(__name__)


@dataclass(frozen=True)
class Form4ClusterEntry:
    symbol: str
    filing_date: dt.date
    as_of_date: dt.date
    total_value: float
    owner_count: int
    transaction_count: int
    event_day_count: int
    max_purchase_pct: float
    median_purchase_pct: float
    weighted_purchase_pct: float
    max_lag_days: int | None = None
    min_lag_days: int | None = None
    has_officer_or_director: bool = False


def _coerce_date(value: Any) -> dt.date | None:
    if value is None:
        return None
    if isinstance(value, dt.datetime):
        return value.date()
    if isinstance(value, dt.date):
        return value
    text = str(value).strip()
    if not text:
        return None
    try:
        return dt.date.fromisoformat(text[:10])
    except ValueError:
        return None


def _coerce_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _coerce_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


class PointInTimeForm4Calendar:
    """Latest-known same-day Form 4 cluster events by filing date."""

    def __init__(self, entries: Iterable[Form4ClusterEntry]) -> None:
        grouped: dict[dt.date, list[Form4ClusterEntry]] = {}
        for entry in entries:
            grouped.setdefault(entry.filing_date, []).append(entry)
        self._entries_by_filing_date = {
            filing_date: tuple(sorted(
                filing_entries,
                key=lambda item: (
                    item.symbol,
                    -item.owner_count,
                    -item.total_value,
                    -item.weighted_purchase_pct,
                ),
            ))
            for filing_date, filing_entries in grouped.items()
        }

    @classmethod
    def from_parquet(cls, path: Path) -> PointInTimeForm4Calendar:
        table = pq.read_table(str(path))
        rows = table.to_pylist()
        entries: list[Form4ClusterEntry] = []
        for row in rows:
            symbol = str(row.get("symbol") or "").strip().upper()
            filing_date = _coerce_date(row.get("filing_date"))
            as_of_date = _coerce_date(row.get("as_of_date") or row.get("filing_date"))
            total_value = _coerce_float(row.get("total_value"))
            owner_count = _coerce_int(row.get("owner_count"))
            transaction_count = _coerce_int(row.get("transaction_count"))
            event_day_count = _coerce_int(row.get("event_day_count"))
            max_purchase_pct = _coerce_float(row.get("max_purchase_pct"))
            median_purchase_pct = _coerce_float(row.get("median_purchase_pct"))
            weighted_purchase_pct = _coerce_float(row.get("weighted_purchase_pct"))
            if (
                not symbol
                or filing_date is None
                or as_of_date is None
                or total_value is None
                or owner_count is None
                or transaction_count is None
                or event_day_count is None
                or max_purchase_pct is None
                or median_purchase_pct is None
                or weighted_purchase_pct is None
            ):
                continue
            entries.append(
                Form4ClusterEntry(
                    symbol=symbol,
                    filing_date=filing_date,
                    as_of_date=as_of_date,
                    total_value=total_value,
                    owner_count=owner_count,
                    transaction_count=transaction_count,
                    event_day_count=event_day_count,
                    max_purchase_pct=max_purchase_pct,
                    median_purchase_pct=median_purchase_pct,
                    weighted_purchase_pct=weighted_purchase_pct,
                    max_lag_days=_coerce_int(row.get("max_lag_days")),
                    min_lag_days=_coerce_int(row.get("min_lag_days")),
                    has_officer_or_director=bool(row.get("has_officer_or_director") or False),
                )
            )
        logger.info(
            "pit_form4_calendar_loaded",
            path=str(path),
            rows=len(entries),
            filing_dates=len({entry.filing_date for entry in entries}),
            symbols=len({entry.symbol for entry in entries}),
        )
        return cls(entries)

    def get_events_between(
        self,
        *,
        start_filing_date: dt.date,
        end_filing_date: dt.date,
        symbols: Iterable[str] | None = None,
    ) -> list[Form4ClusterEntry]:
        symbol_filter = {
            str(symbol).strip().upper()
            for symbol in (symbols or [])
            if str(symbol).strip()
        }
        rows: list[Form4ClusterEntry] = []
        current = start_filing_date
        while current <= end_filing_date:
            for entry in self._entries_by_filing_date.get(current, ()):
                if symbol_filter and entry.symbol not in symbol_filter:
                    continue
                rows.append(entry)
            current += dt.timedelta(days=1)
        return rows


@lru_cache(maxsize=8)
def load_pit_form4_calendar(path_str: str) -> PointInTimeForm4Calendar | None:
    path = Path(path_str)
    if not path.exists():
        logger.info("pit_form4_calendar_missing", path=str(path))
        return None
    return PointInTimeForm4Calendar.from_parquet(path)


__all__ = [
    "Form4ClusterEntry",
    "PointInTimeForm4Calendar",
    "load_pit_form4_calendar",
]