fithia2/libs/export/snapshot_export.py

"""Export feature snapshots + labels to Parquet with train/valid/test split."""
from __future__ import annotations

import datetime as dt
import json
import subprocess
import uuid
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from libs.common.logging import get_logger
from libs.common.time_utils import utc_now

logger = get_logger(__name__)

MANIFEST_FILENAME = "manifest.json"


def _get_git_commit_hash() -> str:
    """Return the current git commit hash (short), or 'unknown'."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return result.stdout.strip() or "unknown"
    except Exception:
        return "unknown"


def _temporal_split(
    rows: list[dict[str, Any]],
    split_policy: str = "temporal_70_15_15",
) -> dict[str, list[dict[str, Any]]]:
    """Split rows into train/valid/test by event_date (temporal order).

    Args:
        rows: List of row dicts that must have an "event_date" field.
        split_policy: E.g. "temporal_70_15_15" → 70% train, 15% valid, 15% test.

    Returns:
        Dict with keys "train", "valid", "test".
    """
    if not rows:
        return {"train": [], "valid": [], "test": []}

    parts = split_policy.replace("temporal_", "").split("_")
    if len(parts) != 3:
        raise ValueError(f"Invalid split_policy: {split_policy}")
    train_pct, valid_pct, _ = (int(p) for p in parts)

    sorted_rows = sorted(rows, key=lambda r: r.get("event_date", ""))
    n = len(sorted_rows)
    n_train = int(n * train_pct / 100)
    n_valid = int(n * valid_pct / 100)

    return {
        "train": sorted_rows[:n_train],
        "valid": sorted_rows[n_train : n_train + n_valid],
        "test": sorted_rows[n_train + n_valid :],
    }


def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table:
    """Convert list of dicts to a PyArrow Table."""
    if not rows:
        return pa.table({})
    # Collect all keys
    keys = list(rows[0].keys())
    arrays: dict[str, list[Any]] = {k: [] for k in keys}
    for row in rows:
        for k in keys:
            arrays[k].append(row.get(k))
    return pa.table({k: pa.array(v) for k, v in arrays.items()})


async def export_dataset_snapshot(
    session: AsyncSession,
    snapshot_id: str | None,
    split_policy: str,
    output_dir: str | Path,
    feature_version: str = "market_v1",
    feature_versions: list[str] | None = None,
    label_version: str = "label-1.0.0",
    parser_version: str = "rule-1.0.0",
    symbols: list[str] | None = None,
) -> dict[str, Any]:
    """Join FeatureSnapshot + EventLabel and export to Parquet.

    Args:
        session: Async DB session.
        snapshot_id: Unique ID for this snapshot (generated if None).
        split_policy: Temporal split policy string (e.g. "temporal_70_15_15").
        output_dir: Root directory for output files.
        feature_version: Snapshot name filter (used when feature_versions is None).
        feature_versions: Merge multiple feature types (e.g. ["market_v1", "event_v1"]).
        label_version: Label version filter for EventLabel.
        parser_version: Parser version filter for Event.

    Returns:
        Manifest dict with metadata and row counts.
    """
    from libs.db.models import Event, EventLabel, FeatureSnapshot, SymbolMaster

    if snapshot_id is None:
        snapshot_id = str(uuid.uuid4())

    versions = feature_versions or [feature_version]

    out_path = Path(output_dir) / snapshot_id
    out_path.mkdir(parents=True, exist_ok=True)

    # Query: JOIN feature_snapshots + event_labels via event_id
    # When merging multiple versions, query all and group by event_id
    stmt = (
        select(FeatureSnapshot, EventLabel, Event)
        .join(EventLabel, FeatureSnapshot.event_id == EventLabel.event_id)
        .join(Event, FeatureSnapshot.event_id == Event.event_id)
        .where(FeatureSnapshot.snapshot_name.in_(versions))
        .where(EventLabel.label_version == label_version)
        .where(EventLabel.label_status.in_(["ok", "truncated"]))
        .where(EventLabel.invalid_event_for_labeling.is_(False))
    )

    if symbols:
        stmt = stmt.where(
            Event.symbol_id.in_(
                select(SymbolMaster.symbol_id).where(
                    SymbolMaster.ticker.in_([s.upper() for s in symbols])
                )
            )
        )

    result = await session.execute(stmt)
    pairs = result.all()

    # Group features by event_id, merge feature_json from all versions
    event_features: dict[str, dict[str, Any]] = {}
    event_labels: dict[str, Any] = {}
    event_dates: dict[str, dt.date] = {}
    for fs, lbl, evt in pairs:
        eid = fs.event_id
        if eid not in event_features:
            event_features[eid] = {}
            event_labels[eid] = lbl
            event_dates[eid] = evt.event_date
        event_features[eid].update(fs.feature_json)

    rows: list[dict[str, Any]] = []
    for eid, features in event_features.items():
        lbl = event_labels[eid]
        row: dict[str, Any] = {
            "event_id": eid,
            "snapshot_name": "+".join(versions),
            "snapshot_version": "1.0.0",
            **features,
            "entry_convention": lbl.entry_convention,
            "reaction_date": lbl.reaction_date.isoformat() if lbl.reaction_date else None,
            "entry_date": lbl.entry_date.isoformat() if lbl.entry_date else None,
            "entry_price": float(lbl.entry_price) if lbl.entry_price else None,
            "fwd_return_1d": float(lbl.fwd_return_1d) if lbl.fwd_return_1d else None,
            "fwd_return_3d": float(lbl.fwd_return_3d) if lbl.fwd_return_3d else None,
            "fwd_return_5d": float(lbl.fwd_return_5d) if lbl.fwd_return_5d else None,
            "hit_pos_1r_within_3d": lbl.hit_pos_1r_within_3d,
            "hit_neg_1r_within_3d": lbl.hit_neg_1r_within_3d,
            "close_up_after_3d": lbl.close_up_after_3d,
            "close_up_after_5d": lbl.close_up_after_5d,
            "mfe_3d": float(lbl.mfe_3d) if lbl.mfe_3d else None,
            "mae_3d": float(lbl.mae_3d) if lbl.mae_3d else None,
            "mfe_5d": float(lbl.mfe_5d) if lbl.mfe_5d else None,
            "mae_5d": float(lbl.mae_5d) if lbl.mae_5d else None,
            "fwd_return_10d": float(lbl.fwd_return_10d) if lbl.fwd_return_10d else None,
            "fwd_return_20d": float(lbl.fwd_return_20d) if lbl.fwd_return_20d else None,
            "mfe_10d": float(lbl.mfe_10d) if lbl.mfe_10d else None,
            "mae_10d": float(lbl.mae_10d) if lbl.mae_10d else None,
            "mfe_20d": float(lbl.mfe_20d) if lbl.mfe_20d else None,
            "mae_20d": float(lbl.mae_20d) if lbl.mae_20d else None,
            "label_status": lbl.label_status,
            "label_version": lbl.label_version,
        }
        # Use DB Event.event_date as authoritative source
        ed = event_dates.get(eid)
        row["event_date"] = ed.isoformat() if ed else row.get("event_date", "")
        rows.append(row)

    logger.info("snapshot_export_rows", snapshot_id=snapshot_id, total=len(rows))

    splits = _temporal_split(rows, split_policy)
    row_counts: dict[str, int] = {}

    for split_name, split_rows in splits.items():
        parquet_path = out_path / f"{split_name}.parquet"
        table = _rows_to_table(split_rows)
        pq.write_table(table, str(parquet_path))
        row_counts[split_name] = len(split_rows)
        logger.info("split_written", split=split_name, rows=len(split_rows), path=str(parquet_path))

    manifest: dict[str, Any] = {
        "snapshot_id": snapshot_id,
        "created_at_utc": utc_now().isoformat(),
        "code_commit_hash": _get_git_commit_hash(),
        "feature_version": "+".join(versions),
        "parser_version": parser_version,
        "label_version": label_version,
        "split_policy": split_policy,
        "row_counts": row_counts,
        "total_rows": len(rows),
        "output_dir": str(out_path),
    }

    manifest_path = out_path / MANIFEST_FILENAME
    manifest_path.write_text(json.dumps(manifest, indent=2))
    logger.info("manifest_written", path=str(manifest_path), snapshot_id=snapshot_id)

    return manifest