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

"""Normalize JoinQuant CSV exports into plugin parquet schemas."""

from __future__ import annotations

import re
from pathlib import Path

import numpy as np
import pandas as pd

from plugins.joinquant.schema import (
    JOINQUANT_FILL_COLUMNS,
    JOINQUANT_PNL_COLUMNS,
    JOINQUANT_POSITION_COLUMNS,
)
from plugins.joinquant.symbols import normalize_symbol_pair, to_joinquant_symbol


def _clean_name(name: object) -> str:
    text = str(name).strip().lower()
    text = re.sub(r"[\s\-.()/]+", "_", text)
    return re.sub(r"[^0-9a-z_]+", "", text).strip("_")


def _clean_columns(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    out.columns = [_clean_name(col) for col in out.columns]
    return out


def _pick(df: pd.DataFrame, candidates: list[str]) -> str | None:
    for candidate in candidates:
        clean = _clean_name(candidate)
        if clean in df.columns:
            return clean
    return None


def _series_or_default(df: pd.DataFrame, candidates: list[str], default: object) -> pd.Series:
    col = _pick(df, candidates)
    if col is None:
        return pd.Series([default] * len(df), index=df.index)
    return df[col]


def _date_series(df: pd.DataFrame) -> pd.Series:
    values = _series_or_default(
        df,
        ["date", "trade_date", "datetime", "time", "created_at"],
        pd.NaT,
    )
    parsed = pd.to_datetime(values, errors="coerce").dt.normalize()
    return parsed.dt.strftime("%Y-%m-%d").fillna("")


def _numeric(values: pd.Series, default: float = 0.0) -> pd.Series:
    return pd.to_numeric(values, errors="coerce").replace([np.inf, -np.inf], np.nan).fillna(default)


def _text(values: pd.Series, default: str = "") -> pd.Series:
    return values.fillna(default).astype(str)


def _portfolio_series(df: pd.DataFrame, portfolio_name: str) -> pd.Series:
    return _text(_series_or_default(df, ["portfolio_name", "portfolio", "strategy"], portfolio_name), portfolio_name)


def _symbol_frame(df: pd.DataFrame) -> pd.DataFrame:
    internal_col = _pick(df, ["symbol_id", "internal_symbol"])
    jq_col = _pick(df, ["jq_symbol", "security", "stock", "symbol", "code", "order_book_id"])

    symbol_ids: list[str] = []
    jq_symbols: list[str] = []
    for idx in df.index:
        internal = df.at[idx, internal_col] if internal_col else None
        jq_value = df.at[idx, jq_col] if jq_col else None
        value = internal if internal is not None and str(internal).strip() else jq_value
        if value is None or not str(value).strip():
            symbol_ids.append("")
            jq_symbols.append("")
            continue
        symbol_id, jq_symbol = normalize_symbol_pair(value)
        if jq_value is not None and str(jq_value).strip():
            try:
                _, jq_symbol = normalize_symbol_pair(jq_value)
            except ValueError:
                jq_symbol = to_joinquant_symbol(symbol_id)
        symbol_ids.append(symbol_id)
        jq_symbols.append(jq_symbol)

    return pd.DataFrame({"symbol_id": symbol_ids, "jq_symbol": jq_symbols}, index=df.index)


def _signed_shares(shares: pd.Series, side: pd.Series) -> pd.Series:
    signed = _numeric(shares, 0.0)
    side_text = side.fillna("").astype(str).str.lower()
    sell = side_text.str.contains("sell|short|close|reduce|-", regex=True)
    buy = side_text.str.contains("buy|long|open|add|\\+", regex=True)
    signed = signed.abs()
    signed = signed.mask(sell, -signed)
    signed = signed.mask(~(sell | buy), _numeric(shares, 0.0))
    return signed


def normalize_fills_csv(path: str | Path, portfolio_name: str) -> pd.DataFrame:
    """Read a JoinQuant fills CSV and return ``JOINQUANT_FILL_COLUMNS``."""
    raw = _clean_columns(pd.read_csv(path))
    symbols = _symbol_frame(raw)
    side = _text(_series_or_default(raw, ["side", "action", "direction"], ""))
    requested = _signed_shares(
        _series_or_default(raw, ["requested_shares", "target_shares", "amount", "order_amount"], 0),
        side,
    )
    filled = _signed_shares(
        _series_or_default(raw, ["filled_shares", "filled", "filled_amount", "deal_amount", "traded_shares"], 0),
        side,
    )
    price = _numeric(_series_or_default(raw, ["fill_price", "price", "avg_cost", "avg_price"], np.nan), np.nan)
    trade_value = _numeric(
        _series_or_default(raw, ["trade_value", "value", "filled_value", "turnover"], np.nan),
        np.nan,
    )
    trade_value = trade_value.fillna((filled * price).abs()).fillna(0.0)

    out = pd.DataFrame({
        "date": _date_series(raw),
        "portfolio_name": _portfolio_series(raw, portfolio_name),
        "symbol_id": symbols["symbol_id"],
        "jq_symbol": symbols["jq_symbol"],
        "order_id": _text(_series_or_default(raw, ["order_id", "id"], "")),
        "side": side,
        "requested_shares": requested.astype(float),
        "filled_shares": filled.astype(float),
        "fill_price": price.astype(float),
        "trade_value": trade_value.astype(float),
        "trade_cost": _numeric(_series_or_default(raw, ["trade_cost", "cost", "commission", "fee"], 0.0), 0.0),
        "blocked": _numeric(_series_or_default(raw, ["blocked", "is_blocked"], 0), 0).astype("int64"),
        "raw_status": _text(_series_or_default(raw, ["raw_status", "status", "order_status"], "")),
    })
    return out[JOINQUANT_FILL_COLUMNS]


def normalize_positions_csv(path: str | Path, portfolio_name: str) -> pd.DataFrame:
    """Read a JoinQuant positions CSV and return ``JOINQUANT_POSITION_COLUMNS``."""
    raw = _clean_columns(pd.read_csv(path))
    symbols = _symbol_frame(raw)
    out = pd.DataFrame({
        "date": _date_series(raw),
        "portfolio_name": _portfolio_series(raw, portfolio_name),
        "symbol_id": symbols["symbol_id"],
        "jq_symbol": symbols["jq_symbol"],
        "position_shares": _numeric(
            _series_or_default(raw, ["position_shares", "shares", "amount", "quantity", "total_amount"], 0),
            0,
        ),
        "position_value": _numeric(
            _series_or_default(raw, ["position_value", "market_value", "value"], 0.0),
            0.0,
        ),
        "cash": _numeric(_series_or_default(raw, ["cash", "available_cash"], np.nan), np.nan),
        "total_value": _numeric(
            _series_or_default(raw, ["total_value", "portfolio_value", "total_asset"], np.nan),
            np.nan,
        ),
    })
    return out[JOINQUANT_POSITION_COLUMNS]


def normalize_pnl_csv(path: str | Path, portfolio_name: str) -> pd.DataFrame:
    """Read a JoinQuant daily PnL CSV and return ``JOINQUANT_PNL_COLUMNS``."""
    raw = _clean_columns(pd.read_csv(path))
    total_value = _numeric(
        _series_or_default(raw, ["total_value", "portfolio_value", "total_asset"], np.nan),
        np.nan,
    )
    pnl = _numeric(_series_or_default(raw, ["pnl", "daily_pnl", "profit", "returns_value"], np.nan), np.nan)
    if pnl.isna().all() and total_value.notna().any():
        pnl = total_value.diff().fillna(0.0)

    out = pd.DataFrame({
        "date": _date_series(raw),
        "portfolio_name": _portfolio_series(raw, portfolio_name),
        "gross_exposure": _numeric(
            _series_or_default(raw, ["gross_exposure", "gross", "positions_value", "market_value"], np.nan),
            np.nan,
        ),
        "net_exposure": _numeric(
            _series_or_default(raw, ["net_exposure", "net"], np.nan),
            np.nan,
        ),
        "cash": _numeric(_series_or_default(raw, ["cash", "available_cash"], np.nan), np.nan),
        "total_value": total_value,
        "pnl": pnl.fillna(0.0),
        "cost": _numeric(_series_or_default(raw, ["cost", "trade_cost", "commission", "fee"], 0.0), 0.0),
        "turnover": _numeric(_series_or_default(raw, ["turnover"], 0.0), 0.0),
    })
    return out[JOINQUANT_PNL_COLUMNS]


def ingest_joinquant_outputs(
    *,
    portfolio_name: str,
    fills_csv: str | Path,
    positions_csv: str | Path,
    pnl_csv: str | Path,
    out_dir: str | Path = "plugins_output/joinquant/ingested",
) -> dict[str, Path]:
    """Normalize JoinQuant CSV exports and write parquet outputs."""
    out_root = Path(out_dir) / portfolio_name
    out_root.mkdir(parents=True, exist_ok=True)

    fills = normalize_fills_csv(fills_csv, portfolio_name)
    positions = normalize_positions_csv(positions_csv, portfolio_name)
    pnl = normalize_pnl_csv(pnl_csv, portfolio_name)

    paths = {
        "fills": out_root / "fills.pq",
        "positions": out_root / "positions.pq",
        "pnl": out_root / "pnl.pq",
    }
    fills.to_parquet(paths["fills"], index=False)
    positions.to_parquet(paths["positions"], index=False)
    pnl.to_parquet(paths["pnl"], index=False)
    return paths