"""Execution simulator: next-open fills under A-share trading constraints.

Execution model (documented convention): a position book targeted from
information available on date ``t`` is executed at ``open[t+1]``. Trades that
violate a :class:`~pipeline.portfolio.constraints.TradeConstraint` (suspension,
price limit, volume cap, …) are clipped; a fully blocked buy leaves the position
at its previous level. Realized PnL marks the *actually filled* book.

The simulator is an ABC + a :class:`ReferenceSimulator`; constraints compose by
intersecting their per-name signed delta bounds.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass

import numpy as np
import pandas as pd

from pipeline.common.schema import FILL_COLUMNS, PNL_COLUMNS
from pipeline.portfolio.constraints import TradeConstraint
from pipeline.portfolio.market_rules import MarketRule, compute_limit_status

logger = logging.getLogger(__name__)


@dataclass
class MarketSlice:
    """Per-name market arrays for one execution date (fixed symbol order)."""

    symbol_ids: np.ndarray
    date: object
    price: np.ndarray          # execution/reference price (the open)
    preclose: np.ndarray
    amount: np.ndarray         # daily turnover value (yuan)
    tradestatus: np.ndarray    # 1 traded / 0 suspended
    is_st: np.ndarray
    limit_status: np.ndarray   # LimitStatus values
    close: np.ndarray          # close, for marking


@dataclass
class TradeContext:
    """Inputs handed to constraints and the fill routine for one date."""

    prev_shares: np.ndarray
    target_shares: np.ndarray
    slice: MarketSlice
    booksize: float


@dataclass
class FillResult:
    """Outcome of executing one date's target against the constraints."""

    realized_shares: np.ndarray
    traded_shares: np.ndarray
    cost: np.ndarray
    blocked: np.ndarray


class ExecutionSimulator(ABC):
    """Abstract execution layer. Subclasses define how a target gets filled."""

    def __init__(self, constraints: list[TradeConstraint] | None = None,
                 cost_bps: float = 0.0, slippage_bps: float = 0.0):
        self.constraints = constraints or []
        self.cost_bps = cost_bps
        self.slippage_bps = slippage_bps

    @abstractmethod
    def fill(self, ctx: TradeContext) -> FillResult:
        """Execute ``ctx.target_shares`` from ``ctx.prev_shares``."""


class ReferenceSimulator(ExecutionSimulator):
    """Reference fill model: clip the desired trade to the composed bounds."""

    def fill(self, ctx: TradeContext) -> FillResult:
        prev = ctx.prev_shares.astype(np.int64)
        target = ctx.target_shares.astype(np.int64)

        # Portfolio-level retargeting hooks (e.g. neutrality), if any.
        for c in self.constraints:
            adjusted = c.adjust_targets(ctx)
            if adjusted is not None:
                target = np.asarray(adjusted, dtype=np.int64)

        desired = target - prev
        n = len(prev)
        low = np.full(n, -np.inf)
        high = np.full(n, np.inf)
        for c in self.constraints:
            lo, hi = c.delta_bounds(ctx)
            low = np.maximum(low, lo)
            high = np.minimum(high, hi)

        # Clip desired delta into the feasible interval; round toward zero so a
        # value/volume cap yields a conservative partial fill.
        clipped = np.clip(desired.astype(np.float64), low, high)
        traded = np.trunc(clipped).astype(np.int64)
        blocked = (traded != desired).astype(np.int64)

        realized = prev + traded
        open_px = np.where(np.isfinite(ctx.slice.price), ctx.slice.price, 0.0)
        trade_value = np.abs(traded.astype(np.float64) * open_px)
        cost = trade_value * (self.cost_bps + self.slippage_bps) / 1e4
        return FillResult(realized, traded, cost, blocked)

    def run(
        self,
        positions_df: pd.DataFrame,
        data_df: pd.DataFrame,
        rule_engine: MarketRule | None = None,
    ) -> tuple[pd.DataFrame, pd.DataFrame]:
        """Simulate the whole book date by date with next-open execution.

        For each signal date ``t`` in ``positions_df`` the target is executed at
        the *next* available data date's open. Returns ``(fills, pnl)`` with
        FILL_COLUMNS / PNL_COLUMNS.

        Args:
            positions_df: POSITION_COLUMNS (uses constructed ``position_shares``).
            data_df: DATA_COLUMNS (open/close/preclose/amount/tradestatus/isST).
            rule_engine: For per-name price-limit bands; default built if None.

        Returns:
            ``(fills_df, pnl_df)``.
        """
        rule_engine = rule_engine or MarketRule()
        portfolio_name = (
            positions_df["portfolio_name"].iloc[0] if len(positions_df) else ""
        )
        # Booksize ≈ the per-date gross dollar target (constant by construction).
        if "target_value" in positions_df.columns and len(positions_df):
            per_date_gross = (positions_df.groupby("date")["target_value"]
                              .apply(lambda s: s.abs().sum()))
            booksize = float(per_date_gross.max()) or 1.0
        else:
            booksize = 1.0

        def wide(df, col):
            return df.pivot_table(index="date", columns="symbol_id",
                                  values=col, aggfunc="first").sort_index()

        tgt = wide(positions_df, "position_shares")
        opn = wide(data_df, "open")
        close = wide(data_df, "close")
        preclose = wide(data_df, "preclose") if "preclose" in data_df.columns else close.shift(1)
        amount = wide(data_df, "amount") if "amount" in data_df.columns else opn * np.inf
        tstat = wide(data_df, "tradestatus") if "tradestatus" in data_df.columns else opn.notna().astype(float)
        st = wide(data_df, "isST") if "isST" in data_df.columns else opn * 0.0

        symbols = sorted(set(tgt.columns) | set(opn.columns))
        tgt = tgt.reindex(columns=symbols)
        opn = opn.reindex(columns=symbols)
        close = close.reindex(columns=symbols)
        preclose = preclose.reindex(columns=symbols)
        amount = amount.reindex(columns=symbols)
        tstat = tstat.reindex(columns=symbols)
        st = st.reindex(columns=symbols)

        sym_arr = np.asarray(symbols, dtype=object)
        n = len(symbols)
        data_dates = list(close.index)
        date_pos = {d: i for i, d in enumerate(data_dates)}

        prev_shares = np.zeros(n, dtype=np.int64)
        mark_prev = None  # last close at which the book was marked
        fill_blocks: list[pd.DataFrame] = []
        pnl_rows: list[dict] = []

        for t in tgt.index:
            # Execute at the next available data date after the signal date t.
            i = date_pos.get(t)
            if i is None or i + 1 >= len(data_dates):
                continue
            e = data_dates[i + 1]

            open_e = opn.loc[e].to_numpy(dtype=np.float64)
            close_e = close.loc[e].to_numpy(dtype=np.float64)
            pre_e = preclose.loc[e].to_numpy(dtype=np.float64)
            amt_e = amount.loc[e].to_numpy(dtype=np.float64)
            tstat_e = np.nan_to_num(tstat.loc[e].to_numpy(dtype=np.float64), nan=0.0)
            st_e = np.nan_to_num(st.loc[e].to_numpy(dtype=np.float64), nan=0.0)
            target = np.nan_to_num(tgt.loc[t].to_numpy(dtype=np.float64), nan=0.0).astype(np.int64)

            _, _, _, limit_pct = rule_engine.get_rules_vectorized(sym_arr, e, st_e)
            limit_status = compute_limit_status(open_e, pre_e, limit_pct)

            mslice = MarketSlice(
                symbol_ids=sym_arr, date=e, price=open_e, preclose=pre_e,
                amount=amt_e, tradestatus=tstat_e, is_st=st_e,
                limit_status=limit_status, close=close_e,
            )
            ctx = TradeContext(prev_shares, target, mslice, booksize)
            res = self.fill(ctx)

            # PnL: overnight gap on the OLD book + intraday on the NEW book - cost.
            if mark_prev is None:
                overnight = 0.0
            else:
                gap = np.where(np.isfinite(open_e) & np.isfinite(mark_prev),
                               open_e - mark_prev, 0.0)
                overnight = float(np.nansum(prev_shares * gap))
            intraday_px = np.where(np.isfinite(close_e) & np.isfinite(open_e),
                                   close_e - open_e, 0.0)
            intraday = float(np.nansum(res.realized_shares * intraday_px))
            cost_total = float(np.nansum(res.cost))
            pnl = overnight + intraday - cost_total

            mark_e = np.where(np.isfinite(close_e), close_e, open_e)
            realized_value = res.realized_shares * np.where(np.isfinite(mark_e), mark_e, 0.0)
            traded_value = np.abs(res.traded_shares * np.where(np.isfinite(open_e), open_e, 0.0))

            nz = res.realized_shares != 0
            fill_blocks.append(pd.DataFrame({
                "symbol_id": symbols,
                "date": e,
                "portfolio_name": portfolio_name,
                "prev_shares": prev_shares,
                "target_shares": target,
                "traded_shares": res.traded_shares,
                "realized_shares": res.realized_shares,
                "blocked": res.blocked,
                "trade_cost": res.cost,
            })[lambda d: (d["traded_shares"] != 0) | (d["realized_shares"] != 0)])

            pnl_rows.append({
                "date": e,
                "portfolio_name": portfolio_name,
                "gross_exposure": float(np.abs(realized_value).sum()),
                "net_exposure": float(realized_value.sum()),
                "pnl": pnl,
                "cost": cost_total,
                "turnover": float(traded_value.sum() / booksize) if booksize else 0.0,
                "n_positions": int(nz.sum()),
            })

            prev_shares = res.realized_shares
            mark_prev = mark_e

        fills_df = (pd.concat(fill_blocks, ignore_index=True)[FILL_COLUMNS]
                    if fill_blocks else pd.DataFrame(columns=FILL_COLUMNS))
        pnl_df = (pd.DataFrame(pnl_rows)[PNL_COLUMNS]
                  if pnl_rows else pd.DataFrame(columns=PNL_COLUMNS))
        logger.info(
            "Simulated '%s': %d exec days, final gross %.0f, total cost %.0f",
            portfolio_name, len(pnl_df),
            pnl_df["gross_exposure"].iloc[-1] if len(pnl_df) else 0.0,
            pnl_df["cost"].sum() if len(pnl_df) else 0.0,
        )
        return fills_df, pnl_df