"""Date-aware A-share market rule engine.

The engine is deliberately separated from alpha/portfolio logic: it answers a
single question — *what lot, increment, sell, and price-limit rules apply to a
given symbol on a given date* — from a data-driven table. New rule changes or
new boards are added by appending rows to :data:`RULE_TABLE`; no branching logic
needs editing.

Boards are detected from the internal ``symbol_id`` prefix (``sh600000`` /
``sz000001`` / ``sh688981`` / ``sz300750``).
"""

from __future__ import annotations

import datetime as _dt
from dataclasses import dataclass
from enum import Enum

import numpy as np


class Board(str, Enum):
    """A-share trading board."""

    MAIN = "main"        # sh60xxxx, sz000/001/002xxx (沪深主板, incl. former SME)
    STAR = "star"        # sh688xxx (科创板 / STAR Market)
    CHINEXT = "chinext"  # sz300xxx (创业板 / ChiNext)
    UNKNOWN = "unknown"


class LimitStatus(int, Enum):
    """Daily price-limit state of a name on a given date.

    Constraints consume this status rather than comparing raw prices, so future
    refinements (一字板 / limit-locked, queue priority, partial fills) only add
    states or richer fill logic without rewriting constraints.
    """

    NORMAL = 0
    UP_LIMIT = 1
    DOWN_LIMIT = -1


@dataclass(frozen=True)
class Rule:
    """Lot/sell/price-limit rule that applies to one (board, date) cell."""

    minimum_open_size: int   # min shares to OPEN (buy) a position
    share_increment: int     # lot granularity above the minimum
    sell_rule: str           # "lot" | "odd_lot_full" (odd residual sellable whole)
    price_limit_pct: float   # daily up/down band as a fraction (e.g. 0.10)


@dataclass(frozen=True)
class RuleSpan:
    """A rule that is valid for ``[valid_from, valid_to)`` on a board."""

    board: Board
    valid_from: _dt.date
    valid_to: _dt.date
    rule: Rule


# --- The rule table ----------------------------------------------------------
# Append rows here for future rule changes or new boards. Order does not matter;
# get_rule selects by board + [valid_from, valid_to) membership.
_MAIN_INCREMENT_CHANGE = _dt.date(2023, 8, 10)
_DATE_MIN = _dt.date(1990, 1, 1)
_DATE_MAX = _dt.date(2999, 12, 31)

#: Price-limit band applied to ST names, overriding the board band.
ST_PRICE_LIMIT_PCT = 0.05

RULE_TABLE: list[RuleSpan] = [
    # 沪深主板: before 2023-08-10 orders must be whole multiples of 100;
    # on/after 2023-08-10 the minimum is still 100 but the increment is 1 share.
    RuleSpan(Board.MAIN, _DATE_MIN, _MAIN_INCREMENT_CHANGE,
             Rule(100, 100, "lot", 0.10)),
    RuleSpan(Board.MAIN, _MAIN_INCREMENT_CHANGE, _DATE_MAX,
             Rule(100, 1, "lot", 0.10)),
    # 科创板 (STAR): from launch, min buy 200, increment 1; a residual holding
    # below 200 (an odd lot) may be sold in full.
    RuleSpan(Board.STAR, _DATE_MIN, _DATE_MAX,
             Rule(200, 1, "odd_lot_full", 0.20)),
    # 创业板 (ChiNext): APPROXIMATION — modeled as post-2023 main-board lots
    # (min 100, increment 1) with a 20% band. Real ChiNext history (e.g. the
    # 2020-08-24 registration-system 20% band, earlier 10% band, 100-share lots)
    # can be added as extra rows here WITHOUT touching get_rule's logic.
    RuleSpan(Board.CHINEXT, _DATE_MIN, _DATE_MAX,
             Rule(100, 1, "lot", 0.20)),
]

#: Fallback for UNKNOWN boards — conservative main-board-like lots, 10% band.
_DEFAULT_RULE = Rule(100, 100, "lot", 0.10)


def detect_board(symbol_id: str) -> Board:
    """Classify a symbol into its trading board from the ``symbol_id`` prefix.

    Args:
        symbol_id: Internal code like ``sh600000`` / ``sz300750``.

    Returns:
        The :class:`Board`; :attr:`Board.UNKNOWN` if no rule matches.
    """
    if len(symbol_id) < 5:
        return Board.UNKNOWN
    exchange, code = symbol_id[:2], symbol_id[2:]
    if exchange == "sh":
        if code.startswith("688"):
            return Board.STAR
        if code.startswith("60"):
            return Board.MAIN
    elif exchange == "sz":
        if code.startswith("300"):
            return Board.CHINEXT
        if code[:3] in ("000", "001", "002"):
            return Board.MAIN
    return Board.UNKNOWN


def _to_date(value) -> _dt.date:
    """Coerce a date / datetime / pandas Timestamp / ISO string to ``date``."""
    if isinstance(value, _dt.datetime):
        return value.date()
    if isinstance(value, _dt.date):
        return value
    # numpy datetime64 / pandas Timestamp / str all accept str() round-trip.
    return _dt.date.fromisoformat(str(value)[:10])


class MarketRule:
    """Resolve lot/sell/price-limit rules for a symbol on a date.

    The engine is stateless; instantiate once and reuse across the run.
    """

    def __init__(self, table: list[RuleSpan] | None = None,
                 default_rule: Rule = _DEFAULT_RULE,
                 st_price_limit_pct: float = ST_PRICE_LIMIT_PCT) -> None:
        self._table = table if table is not None else RULE_TABLE
        self._default = default_rule
        self._st_limit = st_price_limit_pct

    def get_rule(self, symbol_id: str, on, is_st: bool = False) -> Rule:
        """Return the :class:`Rule` for ``symbol_id`` on date ``on``.

        Args:
            symbol_id: Internal symbol code.
            on: Trading date (``date``, ``datetime``, ``Timestamp``, or ISO str).
            is_st: If True, override the price-limit band with the ST band.

        Returns:
            The matching :class:`Rule`, with ST band applied if ``is_st``.
        """
        day = _to_date(on)
        board = detect_board(symbol_id)
        rule = self._default
        for span in self._table:
            if span.board is board and span.valid_from <= day < span.valid_to:
                rule = span.rule
                break
        if is_st:
            rule = Rule(rule.minimum_open_size, rule.share_increment,
                        rule.sell_rule, self._st_limit)
        return rule

    def get_rules_vectorized(
        self, symbol_ids, on, is_st,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Vectorized rule lookup for a whole cross-section on one date.

        Args:
            symbol_ids: Sequence of ``symbol_id`` strings (length N).
            on: The trading date (shared by all names).
            is_st: Boolean array (length N), 1/True where the name is ST.

        Returns:
            Tuple of four numpy arrays, each length N:
            ``(min_open, increment, sell_is_odd_full, limit_pct)`` with dtypes
            ``int64, int64, bool, float64``.
        """
        symbol_ids = np.asarray(symbol_ids, dtype=object)
        is_st = np.asarray(is_st).astype(bool)
        n = len(symbol_ids)
        min_open = np.empty(n, dtype=np.int64)
        increment = np.empty(n, dtype=np.int64)
        odd_full = np.empty(n, dtype=bool)
        limit_pct = np.empty(n, dtype=np.float64)

        # Resolve once per distinct symbol (board only depends on the prefix).
        cache: dict[str, Rule] = {}
        for i, sym in enumerate(symbol_ids):
            rule = cache.get(sym)
            if rule is None:
                rule = self.get_rule(sym, on, is_st=False)
                cache[sym] = rule
            min_open[i] = rule.minimum_open_size
            increment[i] = rule.share_increment
            odd_full[i] = rule.sell_rule == "odd_lot_full"
            limit_pct[i] = self._st_limit if is_st[i] else rule.price_limit_pct
        return min_open, increment, odd_full, limit_pct


def compute_limit_status(
    price, preclose, limit_pct, *, tol: float = 1e-6,
) -> np.ndarray:
    """Classify each name's daily price-limit state.

    A name is at the up (down) limit when its price reaches
    ``preclose * (1 ± limit_pct)`` within ``tol`` relative tolerance.

    Args:
        price: Reference price array (e.g. the open at which we trade).
        preclose: Previous close array.
        limit_pct: Per-name daily band fraction.
        tol: Relative tolerance for the limit comparison.

    Returns:
        ``int8`` array of :class:`LimitStatus` values (length N). Names with
        non-positive or NaN preclose are treated as ``NORMAL``.
    """
    price = np.asarray(price, dtype=np.float64)
    preclose = np.asarray(preclose, dtype=np.float64)
    limit_pct = np.asarray(limit_pct, dtype=np.float64)

    status = np.zeros(price.shape, dtype=np.int8)
    valid = np.isfinite(price) & np.isfinite(preclose) & (preclose > 0)
    up = preclose * (1.0 + limit_pct)
    down = preclose * (1.0 - limit_pct)
    at_up = valid & (price >= up * (1.0 - tol))
    at_down = valid & (price <= down * (1.0 + tol))
    status[at_up] = LimitStatus.UP_LIMIT.value
    status[at_down] = LimitStatus.DOWN_LIMIT.value
    return status