chinese-equity-quant/plugins/joinquant/export_targets.py

"""Export portfolio positions as frozen JoinQuant target files."""

from __future__ import annotations

import hashlib
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Literal

import pandas as pd

from pipeline.common.schema import POSITION_COLUMNS
from plugins.joinquant.schema import JOINQUANT_TARGET_COLUMNS
from plugins.joinquant.symbols import to_joinquant_symbol


ExportMode = Literal["target_shares", "target_value"]


def _date_text(value: object) -> str:
    return pd.Timestamp(value).strftime("%Y-%m-%d")


def _date_file_stem(date_text: str) -> str:
    return pd.Timestamp(date_text).strftime("%Y%m%d")


def _snapshot_root_for(targets_root: Path) -> Path:
    if targets_root.name == "targets":
        return targets_root.parent / "snapshots"
    return targets_root / "snapshots"


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _check_position_columns(df: pd.DataFrame) -> None:
    missing = [col for col in POSITION_COLUMNS if col not in df.columns]
    if missing:
        raise ValueError(f"Positions input missing required columns: {missing}")


def _filter_dates(
    df: pd.DataFrame,
    start_date: str | None,
    end_date: str | None,
    *,
    date_column: str = "date",
) -> pd.DataFrame:
    out = df.copy()
    out[date_column] = pd.to_datetime(out[date_column]).dt.normalize()
    if start_date:
        out = out[out[date_column] >= pd.Timestamp(start_date).normalize()]
    if end_date:
        out = out[out[date_column] <= pd.Timestamp(end_date).normalize()]
    return out


def _read_execution_calendar(path: str | Path) -> pd.DatetimeIndex:
    data = pd.read_parquet(path)
    if "date" not in data.columns:
        raise ValueError("execution calendar parquet must contain a 'date' column")
    dates = pd.to_datetime(data["date"], errors="coerce").dropna().dt.normalize()
    return pd.DatetimeIndex(sorted(dates.unique()))


def _apply_execution_calendar(df: pd.DataFrame, calendar_path: str | Path) -> pd.DataFrame:
    calendar = _read_execution_calendar(calendar_path)
    if calendar.empty:
        raise ValueError("execution calendar contains no dates")

    source_dates = pd.to_datetime(df["date"]).dt.normalize()
    positions = calendar.searchsorted(source_dates, side="right")
    out = df.copy()
    out["source_date"] = source_dates
    out["export_date"] = pd.NaT
    valid = positions < len(calendar)
    out.loc[valid, "export_date"] = calendar.take(positions[valid])
    return out[out["export_date"].notna()].copy()


def build_target_frame(
    positions: pd.DataFrame,
    *,
    portfolio_name: str | None = None,
    mode: ExportMode = "target_shares",
    start_date: str | None = None,
    end_date: str | None = None,
    snapshot_ids: dict[str, str] | None = None,
    execution_calendar_path: str | Path | None = None,
) -> pd.DataFrame:
    """Build normalized JoinQuant target rows from portfolio positions.

    ``target_shares`` is populated from ``position_shares`` because the core
    simulator executes the discretized book, not continuous research shares.
    """
    if mode not in {"target_shares", "target_value"}:
        raise ValueError("mode must be 'target_shares' or 'target_value'")
    _check_position_columns(positions)

    df = positions.copy()
    df["date"] = pd.to_datetime(df["date"]).dt.normalize()
    if execution_calendar_path is not None:
        df = _apply_execution_calendar(df, execution_calendar_path)
        df = df.drop(columns=["date"]).rename(columns={"export_date": "date"})
        df["source_date"] = pd.to_datetime(df["source_date"]).dt.strftime("%Y-%m-%d")
        df = _filter_dates(df, start_date, end_date)
    else:
        df = _filter_dates(df, start_date, end_date)
    if portfolio_name is not None:
        df = df[df["portfolio_name"].astype(str) == portfolio_name]
    if df.empty:
        return pd.DataFrame(columns=JOINQUANT_TARGET_COLUMNS)

    out = pd.DataFrame({
        "date": df["date"].map(_date_text),
        "portfolio_name": df["portfolio_name"].astype(str),
        "symbol_id": df["symbol_id"].astype(str),
        "jq_symbol": df["symbol_id"].map(to_joinquant_symbol),
        "target_shares": pd.to_numeric(df["position_shares"], errors="coerce").fillna(0).astype("int64"),
        "target_value": pd.to_numeric(df["target_value"], errors="coerce").fillna(0.0),
        "target_weight": pd.to_numeric(df["target_weight"], errors="coerce").fillna(0.0),
        "export_mode": mode,
        "snapshot_id": "",
    })

    if snapshot_ids:
        out["snapshot_id"] = out["date"].map(snapshot_ids).fillna("")

    return out[JOINQUANT_TARGET_COLUMNS].sort_values(
        ["date", "portfolio_name", "symbol_id"]
    ).reset_index(drop=True)


def export_targets(
    positions_path: str | Path,
    *,
    portfolio_name: str,
    mode: ExportMode = "target_shares",
    out_dir: str | Path = "plugins_output/joinquant/targets",
    start_date: str | None = None,
    end_date: str | None = None,
    execution_calendar_path: str | Path | None = None,
    force: bool = False,
) -> list[dict[str, object]]:
    """Export one daily CSV/parquet target file plus a snapshot JSON per date.

    Args:
        positions_path: Parquet file produced by ``portfolio build``.
        portfolio_name: Portfolio run to export.
        mode: ``target_shares`` or ``target_value``.
        out_dir: Target root. Files are written to ``out_dir/portfolio_name``.
            If the root is named ``targets``, snapshots are written to the
            sibling ``snapshots`` directory.
        start_date: Optional inclusive start date.
        end_date: Optional inclusive end date.
        execution_calendar_path: Optional daily-bar parquet dataset used to
            shift each position date to the next available execution session.
            This matches the internal simulator's next-open convention.
        force: If false, existing target or snapshot files are treated as
            frozen and cause ``FileExistsError``.

    Returns:
        Snapshot metadata dictionaries, one per exported date.
    """
    positions_path = Path(positions_path)
    targets_root = Path(out_dir)
    snapshot_root = _snapshot_root_for(targets_root)
    targets_portfolio_dir = targets_root / portfolio_name
    snapshots_portfolio_dir = snapshot_root / portfolio_name
    targets_portfolio_dir.mkdir(parents=True, exist_ok=True)
    snapshots_portfolio_dir.mkdir(parents=True, exist_ok=True)

    positions = pd.read_parquet(positions_path)
    filtered = positions.copy()
    filtered["date"] = pd.to_datetime(filtered["date"]).dt.normalize()
    if execution_calendar_path is not None:
        filtered = _apply_execution_calendar(filtered, execution_calendar_path)
        filtered = filtered.drop(columns=["date"]).rename(columns={"export_date": "date"})
        filtered["source_date"] = pd.to_datetime(filtered["source_date"]).dt.strftime("%Y-%m-%d")
        filtered = _filter_dates(filtered, start_date, end_date)
    else:
        filtered = _filter_dates(filtered, start_date, end_date)
    filtered = filtered[filtered["portfolio_name"].astype(str) == portfolio_name]
    if filtered.empty:
        return []

    date_texts = sorted(filtered["date"].map(_date_text).unique())
    snapshot_ids = {
        date_text: f"jq-{portfolio_name}-{date_text}-{uuid.uuid4().hex[:12]}"
        for date_text in date_texts
    }
    targets = build_target_frame(
        filtered,
        portfolio_name=portfolio_name,
        mode=mode,
        snapshot_ids=snapshot_ids,
        execution_calendar_path=None,
    )

    snapshots: list[dict[str, object]] = []
    for date_text, daily in targets.groupby("date", sort=True):
        stem = _date_file_stem(date_text)
        csv_path = targets_portfolio_dir / f"{stem}.csv"
        parquet_path = targets_portfolio_dir / f"{stem}.parquet"
        snapshot_path = snapshots_portfolio_dir / f"{stem}.json"
        existing: Iterable[Path] = (csv_path, parquet_path, snapshot_path)
        if not force:
            conflicts = [str(path) for path in existing if path.exists()]
            if conflicts:
                raise FileExistsError(
                    "Frozen JoinQuant target already exists; use --force to overwrite: "
                    + ", ".join(conflicts)
                )

        daily = daily[JOINQUANT_TARGET_COLUMNS].reset_index(drop=True)
        daily.to_csv(csv_path, index=False)
        daily.to_parquet(parquet_path, index=False)
        file_hash = _sha256_file(csv_path)

        snapshot = {
            "snapshot_id": snapshot_ids[date_text],
            "portfolio_name": portfolio_name,
            "date": date_text,
            "export_mode": mode,
            "source_positions_path": str(positions_path),
            "execution_calendar_path": str(execution_calendar_path) if execution_calendar_path else None,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "n_symbols": int(len(daily)),
            "file_sha256": file_hash,
            "notes": "Frozen JoinQuant target file.",
            "target_csv_path": str(csv_path),
            "target_parquet_path": str(parquet_path),
        }
        snapshot_path.write_text(
            json.dumps(snapshot, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )
        snapshots.append(snapshot)

    return snapshots