data/institutional_flow.py

"""
Institutional investor buy/sell data for Taiwan stocks.

Sources:
  - TWSE T86 daily report (listed stocks)
  - TPEx institutional daily report (OTC/mainboard stocks)

The module is intentionally defensive: public endpoints can be late, absent on
holidays, or temporarily unavailable. Missing rows are represented as neutral
zero-flow records so the prediction pipeline keeps working.
"""

from __future__ import annotations

import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date
from typing import Any

import pandas as pd
import requests
from cachetools import TTLCache

from data.fetcher import detect_exchange

logger = logging.getLogger(__name__)

TWSE_T86_URL = "https://www.twse.com.tw/rwd/zh/fund/T86"
TPEX_DAILY_URL = "https://www.tpex.org.tw/www/zh-tw/insti/dailyTrade"

FLOW_COLUMNS = [
    "foreign_buy",
    "foreign_sell",
    "foreign_net",
    "trust_buy",
    "trust_sell",
    "trust_net",
    "dealer_buy",
    "dealer_sell",
    "dealer_net",
    "institutional_buy",
    "institutional_sell",
    "institutional_net",
]

_DAILY_CACHE: TTLCache = TTLCache(maxsize=600, ttl=6 * 60 * 60)
_FLOW_CACHE: TTLCache = TTLCache(maxsize=100, ttl=30 * 60)


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


def _institutional_fetch_warnings_enabled() -> bool:
    value = os.getenv("ENABLE_INSTITUTIONAL_FETCH_WARNINGS", "0").strip().lower()
    return value in {"1", "true", "yes", "on"}


def _default_max_days() -> int:
    configured = os.getenv("INSTITUTIONAL_FLOW_DAYS")
    if configured is not None:
        try:
            return max(0, int(configured))
        except ValueError:
            logger.warning("Invalid INSTITUTIONAL_FLOW_DAYS=%r; falling back to default", configured)

    # Hugging Face free CPU and LIGHTWEIGHT_MODE should avoid hundreds of
    # external requests on cold start. 60 trading days still supports the 20d
    # institutional z-score and 5d flow features.
    if os.getenv("LIGHTWEIGHT_MODE", "0") == "1" or os.getenv("SPACE_ID"):
        return 60
    return 260


def _default_workers() -> int:
    configured = os.getenv("INSTITUTIONAL_FLOW_WORKERS")
    if configured is not None:
        try:
            return max(1, int(configured))
        except ValueError:
            logger.warning("Invalid INSTITUTIONAL_FLOW_WORKERS=%r; falling back to default", configured)
    return 2 if (os.getenv("LIGHTWEIGHT_MODE", "0") == "1" or os.getenv("SPACE_ID")) else 6


def _parse_int(value: Any) -> int:
    """Parse TWSE/TPEx comma-formatted integer fields."""
    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 _safe_get(row: list[Any], idx: int) -> int:
    return _parse_int(row[idx]) if idx < len(row) else 0


def _neutral_row(stock_no: str, date_str: str, source: str = "") -> dict:
    row = {
        "date": date_str,
        "stock_no": stock_no,
        "institutional_available": False,
        "institutional_source": source,
        "institutional_as_of": None,
        "institutional_carry_forward": False,
    }
    row.update({col: 0 for col in FLOW_COLUMNS})
    return row


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


def parse_twse_t86(payload: dict[str, Any], date_str: str) -> dict[str, dict]:
    """
    Parse TWSE T86 JSON into a stock_no -> flow-row mapping.

    TWSE columns:
      2-4 foreign ex-dealer buy/sell/net
      8-10 investment trust buy/sell/net
      11 dealer total net
      12-17 dealer self/hedge buy/sell/net
      18 three-institution net
    """
    if payload.get("stat") != "OK":
        return {}

    rows: dict[str, dict] = {}
    for raw in payload.get("data", []) or []:
        if len(raw) < 2:
            continue
        code = _normalize_code(str(raw[0]))

        foreign_buy = _safe_get(raw, 2)
        foreign_sell = _safe_get(raw, 3)
        foreign_net = _safe_get(raw, 4)

        trust_buy = _safe_get(raw, 8)
        trust_sell = _safe_get(raw, 9)
        trust_net = _safe_get(raw, 10)

        dealer_net = _safe_get(raw, 11)
        dealer_self_buy = _safe_get(raw, 12)
        dealer_self_sell = _safe_get(raw, 13)
        dealer_hedge_buy = _safe_get(raw, 15)
        dealer_hedge_sell = _safe_get(raw, 16)
        dealer_buy = dealer_self_buy + dealer_hedge_buy
        dealer_sell = dealer_self_sell + dealer_hedge_sell

        institutional_net = _safe_get(raw, 18)
        institutional_buy = foreign_buy + trust_buy + dealer_buy
        institutional_sell = foreign_sell + trust_sell + dealer_sell

        rows[code] = {
            "date": date_str,
            "stock_no": code,
            "foreign_buy": foreign_buy,
            "foreign_sell": foreign_sell,
            "foreign_net": foreign_net,
            "trust_buy": trust_buy,
            "trust_sell": trust_sell,
            "trust_net": trust_net,
            "dealer_buy": dealer_buy,
            "dealer_sell": dealer_sell,
            "dealer_net": dealer_net,
            "institutional_buy": institutional_buy,
            "institutional_sell": institutional_sell,
            "institutional_net": institutional_net,
            "institutional_available": True,
            "institutional_source": "TWSE_T86",
            "institutional_as_of": date_str,
            "institutional_carry_forward": False,
        }
    return rows


def parse_tpex_daily(payload: dict[str, Any], date_str: str) -> dict[str, dict]:
    """
    Parse TPEx institutional daily JSON into a stock_no -> flow-row mapping.

    TPEx repeats generic buy/sell/net field names by investor group. The stable
    positional layout is:
      2-4 foreign ex-dealer, 11-13 investment trust, 20-22 dealer total,
      23 three-institution net.
    """
    tables = payload.get("tables") or []
    if not tables:
        return {}

    data = tables[0].get("data", []) or []
    rows: dict[str, dict] = {}
    for raw in data:
        if len(raw) < 2:
            continue
        code = _normalize_code(str(raw[0]))

        foreign_buy = _safe_get(raw, 2)
        foreign_sell = _safe_get(raw, 3)
        foreign_net = _safe_get(raw, 4)

        trust_buy = _safe_get(raw, 11)
        trust_sell = _safe_get(raw, 12)
        trust_net = _safe_get(raw, 13)

        dealer_buy = _safe_get(raw, 20)
        dealer_sell = _safe_get(raw, 21)
        dealer_net = _safe_get(raw, 22)

        institutional_net = _safe_get(raw, 23)
        institutional_buy = foreign_buy + trust_buy + dealer_buy
        institutional_sell = foreign_sell + trust_sell + dealer_sell

        rows[code] = {
            "date": date_str,
            "stock_no": code,
            "foreign_buy": foreign_buy,
            "foreign_sell": foreign_sell,
            "foreign_net": foreign_net,
            "trust_buy": trust_buy,
            "trust_sell": trust_sell,
            "trust_net": trust_net,
            "dealer_buy": dealer_buy,
            "dealer_sell": dealer_sell,
            "dealer_net": dealer_net,
            "institutional_buy": institutional_buy,
            "institutional_sell": institutional_sell,
            "institutional_net": institutional_net,
            "institutional_available": True,
            "institutional_source": "TPEX_DAILY",
            "institutional_as_of": date_str,
            "institutional_carry_forward": False,
        }
    return rows


def _fetch_daily_market(source: str, date_str: str) -> dict[str, dict]:
    """Fetch and parse all institutional rows for one market/date."""
    cache_key = f"{source}:{date_str}"
    cached = _DAILY_CACHE.get(cache_key)
    if cached is not None:
        return cached

    headers = {"User-Agent": "Mozilla/5.0"}
    try:
        if source == "TWSE":
            resp = requests.get(
                TWSE_T86_URL,
                params={
                    "response": "json",
                    "date": date_str.replace("-", ""),
                    "selectType": "ALLBUT0999",
                },
                headers=headers,
                timeout=8,
            )
            resp.raise_for_status()
            rows = parse_twse_t86(resp.json(), date_str)
        else:
            resp = requests.get(
                TPEX_DAILY_URL,
                params={
                    "date": date_str.replace("-", "/"),
                    "type": "Daily",
                    "response": "json",
                },
                headers=headers,
                timeout=8,
            )
            resp.raise_for_status()
            rows = parse_tpex_daily(resp.json(), date_str)
    except Exception as exc:
        log = logger.warning if _institutional_fetch_warnings_enabled() else logger.debug
        log("institutional %s fetch failed for %s: %s", source, date_str, exc)
        rows = {}

    _DAILY_CACHE[cache_key] = rows
    return rows


def _row_for_date(stock_no: str, date_str: str, primary_exchange: str) -> dict:
    """Return one stock's institutional row for a date, trying both markets."""
    bare = _normalize_code(stock_no)
    sources = ["TPEX", "TWSE"] if primary_exchange == "TPEX" else ["TWSE", "TPEX"]

    for source in sources:
        daily = _fetch_daily_market(source, date_str)
        if bare in daily:
            return daily[bare]
    return _neutral_row(bare, date_str)


def fetch_institutional_flow(
    stock_no: str,
    dates: list[str] | pd.Series | pd.Index,
    *,
    exchange: str | None = None,
    max_days: int | None = None,
) -> pd.DataFrame:
    """
    Fetch institutional flow rows aligned to the given price dates.

    Only the most recent max_days are fetched from the public endpoints. Older
    rows are neutral to keep initial training bounded and predictable.
    """
    bare = _normalize_code(stock_no)
    if max_days is None:
        max_days = _default_max_days()

    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()
    if not _institutional_flow_enabled() or max_days <= 0:
        return pd.DataFrame([_neutral_row(bare, d) for d in date_strings])

    selected = date_strings[-max_days:]
    cache_key = f"{bare}:{exchange or ''}:{','.join(selected)}"
    cached = _FLOW_CACHE.get(cache_key)
    if cached is not None:
        return cached.copy()

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

    workers = min(8, _default_workers())
    fetched_by_date: dict[str, dict] = {}
    with ThreadPoolExecutor(max_workers=workers) as executor:
        future_map = {
            executor.submit(_row_for_date, bare, d, primary_exchange): d
            for d in selected
        }
        for future in as_completed(future_map):
            d = future_map[future]
            try:
                fetched_by_date[d] = future.result()
            except Exception as exc:
                log = logger.warning if _institutional_fetch_warnings_enabled() else logger.debug
                log("institutional row failed for %s %s: %s", bare, d, exc)
                fetched_by_date[d] = _neutral_row(bare, d)

    rows.extend(fetched_by_date.get(d, _neutral_row(bare, d)) for d in selected)
    df = pd.DataFrame(rows).sort_values("date").reset_index(drop=True)

    # If the latest trading day is not published yet, use the most recent
    # already-published row for today's prediction without leaking future data.
    if not df.empty and not bool(df.iloc[-1].get("institutional_available", False)):
        prior = df[df["institutional_available"] == True]  # noqa: E712
        if not prior.empty:
            prior_row = prior.iloc[-1]
            last_idx = df.index[-1]
            for col in FLOW_COLUMNS:
                df.at[last_idx, col] = prior_row[col]
            df.at[last_idx, "institutional_as_of"] = prior_row["institutional_as_of"]
            df.at[last_idx, "institutional_source"] = prior_row["institutional_source"]
            df.at[last_idx, "institutional_carry_forward"] = True

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


def add_institutional_flow(
    df: pd.DataFrame,
    stock_no: str,
    *,
    exchange: str | None = None,
    max_days: int | None = None,
) -> pd.DataFrame:
    """Merge institutional flow columns into an OHLCV/indicator DataFrame."""
    if df.empty or "date" not in df.columns:
        return df

    out = df.copy()
    if not _institutional_flow_enabled():
        for col in FLOW_COLUMNS:
            out[col] = 0.0
        out["institutional_available"] = False
        out["institutional_source"] = ""
        out["institutional_as_of"] = None
        out["institutional_carry_forward"] = False
        return out

    flow = fetch_institutional_flow(
        stock_no,
        out["date"],
        exchange=exchange,
        max_days=max_days,
    )
    if flow.empty:
        for col in FLOW_COLUMNS:
            out[col] = 0
        out["institutional_available"] = False
        out["institutional_source"] = ""
        out["institutional_as_of"] = None
        out["institutional_carry_forward"] = False
        return out

    out["_flow_date"] = pd.to_datetime(out["date"]).dt.strftime("%Y-%m-%d")
    merged = out.merge(
        flow,
        how="left",
        left_on="_flow_date",
        right_on="date",
        suffixes=("", "_flow"),
    )
    merged = merged.drop(columns=[c for c in ["_flow_date", "date_flow", "stock_no_flow"] if c in merged.columns])

    for col in FLOW_COLUMNS:
        merged[col] = pd.to_numeric(merged.get(col, 0), errors="coerce").fillna(0).astype(float)
    merged["institutional_available"] = merged.get("institutional_available", False).fillna(False).astype(bool)
    merged["institutional_source"] = merged.get("institutional_source", "").fillna("")
    merged["institutional_as_of"] = merged.get("institutional_as_of", None)
    merged["institutional_carry_forward"] = merged.get("institutional_carry_forward", False).fillna(False).astype(bool)
    return merged