data/margin_flow.py

"""
Taiwan margin trading (融資融券) data fetcher.

Source: TWSE MI_MARGN API (TWSE-listed stocks only; OTC has no equivalent public endpoint).

Derived ML features (added via add_margin_flow):
  margin_balance        raw daily margin balance (trading units = 1000 shares each)
  short_balance         raw daily short balance (trading units)
  short_margin_ratio    short_balance / margin_balance — 券資比 (OTC-style pressure gauge)
  margin_5d_change_pct  5-day % change in margin_balance (rising = retail adding leverage)
  margin_buy_pressure   margin_buy / (margin_buy + margin_sell) — direction of margin flow
  chip_squeeze_signal   foreign_net > 0 AND margin_5d_change_pct > 0.05 (squeeze setup)
"""

from __future__ import annotations

import logging
import os
from datetime import date, timedelta

import pandas as pd
import requests
from cachetools import TTLCache

logger = logging.getLogger(__name__)

TWSE_MARGN_URL = "https://www.twse.com.tw/rwd/zh/marginTrading/MI_MARGN"

MARGIN_COLUMNS = [
    "margin_balance",
    "short_balance",
    "margin_buy",
    "margin_sell",
    "margin_prev_balance",
    "short_prev_balance",
]

# 6h TTL — published once per day, no point refreshing more often
_DAILY_CACHE: TTLCache = TTLCache(maxsize=400, ttl=6 * 60 * 60)
_FLOW_CACHE: TTLCache = TTLCache(maxsize=100, ttl=30 * 60)

_HEADERS = {"User-Agent": "Mozilla/5.0"}


def _margin_enabled() -> bool:
    return os.getenv("ENABLE_MARGIN_FLOW", "1").strip().lower() not in {"0", "false", "no", "off"}


def _margin_fetch_warnings_enabled() -> bool:
    return os.getenv("ENABLE_MARGIN_FETCH_WARNINGS", "0").strip().lower() in {"1", "true", "yes", "on"}


def _parse_int(value) -> int:
    if value is None:
        return 0
    text = str(value).strip().replace(",", "")
    if text in ("", "--", "-", " "):
        return 0
    try:
        return int(float(text))
    except (TypeError, ValueError):
        return 0


def _neutral_row(stock_no: str, date_str: str) -> dict:
    row = {"date": date_str, "stock_no": stock_no, "margin_available": False}
    row.update({col: 0 for col in MARGIN_COLUMNS})
    row["short_buy"] = 0
    row["short_sell"] = 0
    return row


def _normalize_code(stock_no: str) -> str:
    return stock_no.replace(".TW", "").replace(".TWO", "").strip()


def parse_twse_margn(payload: dict, date_str: str) -> dict[str, dict]:
    """
    Parse TWSE MI_MARGN JSON.

    tables[1] = 融資融券彙總 with columns:
      0  代號  1  名稱
      2  融資買進  3  融資賣出  4  現金償還  5  前日餘額  6  今日餘額  7  次日限額
      8  融券買進  9  融券賣出  10 現券償還  11 前日餘額  12 今日餘額  13 次日限額
      14 資券互抵  15 註記
    """
    tables = payload.get("tables") or []
    if len(tables) < 2:
        return {}

    rows: dict[str, dict] = {}
    for raw in tables[1].get("data", []) or []:
        if len(raw) < 13:
            continue
        code = _normalize_code(str(raw[0]))
        rows[code] = {
            "date": date_str,
            "stock_no": code,
            "margin_buy": _parse_int(raw[2]),
            "margin_sell": _parse_int(raw[3]),
            "margin_prev_balance": _parse_int(raw[5]),
            "margin_balance": _parse_int(raw[6]),
            "short_buy": _parse_int(raw[8]),
            "short_sell": _parse_int(raw[9]),
            "short_prev_balance": _parse_int(raw[11]),
            "short_balance": _parse_int(raw[12]),
            "margin_available": True,
        }
    return rows


def _fetch_daily(date_str: str) -> dict[str, dict]:
    """Fetch and cache all margin rows for one date."""
    cached = _DAILY_CACHE.get(date_str)
    if cached is not None:
        return cached

    try:
        resp = requests.get(
            TWSE_MARGN_URL,
            params={"response": "json", "date": date_str.replace("-", ""), "selectType": "ALL"},
            headers=_HEADERS,
            timeout=10,
        )
        resp.raise_for_status()
        rows = parse_twse_margn(resp.json(), date_str)
    except Exception as exc:
        log = logger.warning if _margin_fetch_warnings_enabled() else logger.debug
        log("margin fetch failed for %s: %s", date_str, exc)
        rows = {}

    _DAILY_CACHE[date_str] = rows
    return rows


def _find_nearest_published(date_str: str, lookback: int = 5) -> dict[str, dict]:
    """Walk back up to `lookback` trading days to find a published margin dataset."""
    dt = date.fromisoformat(date_str)
    for _ in range(lookback):
        rows = _fetch_daily(dt.strftime("%Y-%m-%d"))
        if rows:
            return rows
        dt -= timedelta(days=1)
    return {}


def fetch_margin_flow(
    stock_no: str,
    dates: list[str] | pd.Series | pd.Index,
    *,
    max_days: int = 260,
) -> pd.DataFrame:
    """
    Return daily margin rows aligned to `dates`.
    OTC stocks (TPEX) get neutral rows — TWSE doesn't publish their margin data.
    """
    bare = _normalize_code(stock_no)
    if not _margin_enabled() or max_days <= 0:
        return pd.DataFrame([_neutral_row(bare, str(d)) for d in dates])

    date_strings = [pd.to_datetime(d).strftime("%Y-%m-%d") for d in list(dates) if pd.notna(d)]
    if not date_strings:
        return pd.DataFrame()

    cache_key = f"margin:{bare}:{date_strings[-1]}:{len(date_strings)}"
    cached = _FLOW_CACHE.get(cache_key)
    if cached is not None:
        return cached.copy()

    selected = date_strings[-max_days:]
    older = date_strings[: max(0, len(date_strings) - len(selected))]
    rows = [_neutral_row(bare, d) for d in older]

    for date_str in selected:
        daily = _fetch_daily(date_str)
        if bare in daily:
            rows.append(daily[bare])
        else:
            rows.append(_neutral_row(bare, date_str))

    df = pd.DataFrame(rows).sort_values("date").reset_index(drop=True)

    # Carry-forward: if latest day not yet published, use most recent published row
    if not df.empty and not bool(df.iloc[-1].get("margin_available", False)):
        prior = df[df["margin_available"] == True]  # noqa: E712
        if not prior.empty:
            prior_row = prior.iloc[-1]
            last_idx = df.index[-1]
            for col in MARGIN_COLUMNS:
                df.at[last_idx, col] = prior_row[col]
            df.at[last_idx, "margin_available"] = True

    _FLOW_CACHE[cache_key] = df.copy()
    return df


def add_margin_flow(df: pd.DataFrame, stock_no: str, *, max_days: int = 260) -> pd.DataFrame:
    """Merge margin columns into an OHLCV/indicator DataFrame."""
    if df.empty or "date" not in df.columns:
        return df

    out = df.copy()

    # OTC stocks: no TWSE margin data, return neutral zeros
    from data.fetcher import detect_exchange
    if detect_exchange(_normalize_code(stock_no)) == "TPEX":
        for col in MARGIN_COLUMNS + ["short_buy", "short_sell"]:
            out[col] = 0.0
        out["margin_available"] = False
        return out

    flow = fetch_margin_flow(stock_no, out["date"], max_days=max_days)
    if flow.empty:
        for col in MARGIN_COLUMNS + ["short_buy", "short_sell"]:
            out[col] = 0.0
        out["margin_available"] = False
        return out

    out["_mdate"] = pd.to_datetime(out["date"]).dt.strftime("%Y-%m-%d")
    flow_cols = ["date", "margin_available"] + MARGIN_COLUMNS + ["short_buy", "short_sell"]
    for c in ["short_buy", "short_sell"]:
        if c not in flow.columns:
            flow[c] = 0
    merged = out.merge(
        flow[[c for c in flow_cols if c in flow.columns]],
        how="left",
        left_on="_mdate",
        right_on="date",
        suffixes=("", "_mflow"),
    )
    merged = merged.drop(columns=["_mdate"] + [c for c in ["date_mflow"] if c in merged.columns])
    for col in MARGIN_COLUMNS + ["short_buy", "short_sell"]:
        merged[col] = pd.to_numeric(merged.get(col, 0), errors="coerce").fillna(0).astype(float)
    merged["margin_available"] = merged.get("margin_available", False).fillna(False).astype(bool)
    return merged