"""
Event-driven long-only portfolio simulator for factor research.

Differences from ``main.backtest`` (the engine behind the web UI numbers):

* Execution lag. A score computed from the close of day s is traded at the
  OPEN of day s+1 (``exec_at="open"``) or the CLOSE of day s+1
  (``exec_at="close"``). ``main.backtest`` fills at the same close the signal
  saw, which the live trader cannot do.
* Holdings drift between rebalances; costs are charged on actual trades.
* Tradability. A stock that is suspended, or opens at limit-up (buy) or
  limit-down (sell) on the execution day keeps its current position. Stocks
  outside the point-in-time index on the signal date are never bought.
* Costs are asymmetric (CN stamp tax on sells) and include an optional fixed
  per-order fee expressed against a reference capital.

Conventions: every DataFrame shares one date index and one column set
(``Panel.close.columns``). ``score`` is "higher is better"; NaN means
"not eligible".
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import pandas as pd


@dataclass
class Costs:
    buy: float = 0.0010             # fraction of traded value
    sell: float = 0.0010
    sell_tax: pd.Series | None = None  # extra sell-side tax by date (CN stamp duty)
    fixed_fee: float = 0.0          # currency per order
    capital: float = 10_000.0       # reference capital for the fixed fee


@dataclass
class Panel:
    market: str
    open: pd.DataFrame
    high: pd.DataFrame
    low: pd.DataFrame
    close: pd.DataFrame
    volume: pd.DataFrame
    member: pd.DataFrame                 # bool: in PIT index on that date
    can_buy: pd.DataFrame                # bool: buyable at that day's open
    can_sell: pd.DataFrame               # bool: sellable at that day's open
    cash_ret: pd.Series                  # daily return of the idle-cash asset
    benchmark: pd.Series                 # index level (cap-weighted)
    costs: Costs = field(default_factory=Costs)
    extra: dict = field(default_factory=dict)   # market-specific fields (turn, pe, ...)

    @property
    def dates(self) -> pd.DatetimeIndex:
        return self.close.index

    def slice(self, start=None, end=None) -> "Panel":
        idx = self.dates
        m = np.ones(len(idx), bool)
        if start is not None:
            m &= idx >= pd.Timestamp(start)
        if end is not None:
            m &= idx <= pd.Timestamp(end)
        def sub(x):
            if isinstance(x, (pd.Series, pd.DataFrame)) and x.index.equals(idx):
                return x.loc[m]
            return x
        return Panel(self.market, sub(self.open), sub(self.high), sub(self.low), sub(self.close),
                     sub(self.volume), sub(self.member), sub(self.can_buy), sub(self.can_sell),
                     sub(self.cash_ret), sub(self.benchmark), self.costs,
                     {k: sub(v) for k, v in self.extra.items() if not k.startswith("_")})


@dataclass
class Result:
    equity: pd.Series
    returns: pd.Series
    turnover: pd.Series        # one-way traded value / equity, per day
    holdings: pd.DataFrame     # weights after each rebalance (rows = exec dates)
    exposure: pd.Series        # invested fraction at each close
    name: str = ""


def rebalance_dates(dates: pd.DatetimeIndex, every: int | str, offset: int = 0,
                    start: int = 0) -> np.ndarray:
    """Boolean mask of signal dates.

    ``every`` is a trading-day interval (5 = weekly) or "W"/"M" for the last
    trading day of each calendar week/month.
    """
    mask = np.zeros(len(dates), bool)
    if isinstance(every, str):
        key = dates.to_period(every[0])
        last = pd.Series(np.arange(len(dates)), index=dates).groupby(key).max().values
        mask[last] = True
        mask[:start] = False
        return mask
    mask[start + offset::every] = True
    return mask


def select_top(score_row: np.ndarray, top_n: int, held: np.ndarray | None = None,
               buffer: float = 0.0) -> np.ndarray:
    """Indices of the chosen names. With a buffer, current holdings stay while
    their rank is within ``top_n * (1 + buffer)``."""
    valid = np.where(~np.isnan(score_row))[0]
    if len(valid) == 0:
        return valid
    order = valid[np.argsort(-score_row[valid], kind="stable")]
    if buffer <= 0 or held is None or not held.any():
        return order[:top_n]
    rank = np.empty(len(score_row)); rank[:] = np.inf
    rank[order] = np.arange(len(order))
    keep = [i for i in np.where(held)[0] if rank[i] < top_n * (1 + buffer)]
    keep = sorted(keep, key=lambda i: rank[i])[:top_n]
    chosen = list(keep)
    for i in order:
        if len(chosen) >= top_n:
            break
        if i not in keep:
            chosen.append(i)
    return np.array(chosen, int)


def zvol_weights(score_row: np.ndarray, idx: np.ndarray, vol: np.ndarray,
                 lo: float = 0.5, hi: float = 2.0) -> np.ndarray:
    """Grinold (1994) alpha = vol x IC x z, so mean-variance weights with a
    diagonal risk model are proportional to z / vol. z is the normal score of
    the name's cross-sectional percentile; weights are clipped to
    [lo/N, hi/N] and renormalized."""
    from scipy.stats import norm
    valid = ~np.isnan(score_row)
    pct = np.full(len(score_row), np.nan)
    order = np.argsort(score_row[valid], kind="stable")
    ranks = np.empty(valid.sum()); ranks[order] = (np.arange(valid.sum()) + 0.5) / valid.sum()
    pct[valid] = ranks
    z = norm.ppf(pct[idx])
    v = np.where(vol > 0, vol, np.nan)
    v = np.nan_to_num(v, nan=np.nanmedian(v) if np.isfinite(v).any() else 1.0)
    raw = np.clip(z, 0.01, None) / v
    n = len(idx)
    w = raw / raw.sum()
    for _ in range(5):
        w = np.clip(w, lo / n, hi / n)
        w = w / w.sum()
    return w


def simulate(panel: Panel, score: pd.DataFrame | None = None, top_n: int = 10,
             every: int | str = 5, offset: int = 0, warmup: int = 0,
             weighting: str = "equal", buffer: float = 0.0,
             exposure: pd.Series | None = None, exec_at: str = "open",
             target_weights: pd.DataFrame | None = None, min_trade: float = 0.002,
             dyn_n: tuple[int, float] | None = None, name: str = "") -> Result:
    """Run a long-only top-N strategy.

    score           : higher = better, evaluated at each signal date's close.
    target_weights  : alternative to score — explicit weights per signal date
                      (rows need only exist on signal dates; NaN rows skip).
    weighting       : "equal" | "invvol" (60-day vol) | "score" (score-proportional, clipped >0)
                      | "zvol" (normal score / 60-day vol, see ``zvol_weights``)
    dyn_n           : (min_n, threshold): hold every member whose score is at least
                      ``threshold`` (e.g. a mean factor percentile of 0.8), at least
                      ``min_n`` and at most ``top_n`` names.
    exposure        : fraction of equity to invest, read on the signal date
                      (e.g. a market-trend or vol-target overlay); rest earns cash_ret.
    exec_at         : "open" → trade at next day's open; "close" → next day's close.
    min_trade       : resizing trades smaller than this fraction of equity are
                      skipped (the live trader skips tiny orders too); full
                      entries and exits always trade.
    """
    dates = panel.dates
    n_days, n_assets = panel.close.shape
    C = panel.close.ffill().values
    O = panel.open.values
    O = np.where(np.isnan(O), C, O)          # missing open → fall back to close
    prevC = np.vstack([C[:1], C[:-1]])
    member = panel.member.values
    can_buy = panel.can_buy.values
    can_sell = panel.can_sell.values
    cash_ret = panel.cash_ret.reindex(dates).fillna(0.0).values
    sell_tax = (panel.costs.sell_tax.reindex(dates).fillna(0.0).values
                if panel.costs.sell_tax is not None else np.zeros(n_days))
    S = score.reindex(index=dates, columns=panel.close.columns).values if score is not None else None
    TW = (target_weights.reindex(index=dates, columns=panel.close.columns).values
          if target_weights is not None else None)
    expo = exposure.reindex(dates).ffill().fillna(1.0).values if exposure is not None else np.ones(n_days)
    if weighting in ("invvol", "zvol"):
        rets = panel.close.pct_change(fill_method=None)
        vol = rets.rolling(60, min_periods=40).std().values
    sig_mask = rebalance_dates(dates, every, offset, start=warmup)

    h = np.zeros(n_assets)     # value held per asset (at last mark)
    cash = 1.0
    eq = np.empty(n_days); turn = np.zeros(n_days); inv = np.zeros(n_days)
    hold_rows, hold_idx = [], []
    pending = None             # target weight vector awaiting execution

    for d in range(n_days):
        # mark overnight / full-day moves
        with np.errstate(invalid="ignore", divide="ignore"):
            gap = np.where(prevC[d] > 0, O[d] / prevC[d], 1.0)
            intraday = np.where(O[d] > 0, C[d] / O[d], 1.0)
        gap = np.nan_to_num(gap, nan=1.0, posinf=1.0)
        intraday = np.nan_to_num(intraday, nan=1.0, posinf=1.0)
        cash *= 1 + cash_ret[d]

        if pending is not None:
            if exec_at == "open":
                h *= gap
            else:
                h *= gap * intraday
            V = h.sum() + cash
            tgt = pending * V
            # untradable names keep their current value
            frozen_buy = (tgt > h) & ~can_buy[d]
            frozen_sell = (tgt < h) & ~can_sell[d]
            frozen = frozen_buy | frozen_sell
            if frozen.any():
                tgt[frozen] = h[frozen]
                free = ~frozen & (pending > 0)
                room = V - tgt[frozen].sum()
                if free.any() and pending[free].sum() > 0:
                    budget = min(room, pending[free].sum() * V)
                    tgt[free] = pending[free] / pending[free].sum() * max(budget, 0.0)
            if min_trade > 0:
                tiny = (np.abs(tgt - h) < min_trade * V) & (tgt > 0) & (h > 0)
                tgt[tiny] = h[tiny]
            trade = tgt - h
            buys = trade[trade > 0].sum()
            sells = -trade[trade < 0].sum()
            n_orders = int((np.abs(trade) > 1e-4 * V).sum())
            cost = buys * panel.costs.buy + sells * (panel.costs.sell + sell_tax[d])
            # Equity is normalized to 1.0 == costs.capital, so a fixed fee per
            # order is fixed_fee / capital in these units at any book size.
            if panel.costs.fixed_fee > 0:
                cost += n_orders * panel.costs.fixed_fee / panel.costs.capital
            cash += h.sum() - tgt.sum() - cost
            h = tgt
            turn[d] = (buys + sells) / 2 / V if V > 0 else 0.0
            hold_rows.append(h / max(h.sum() + cash, 1e-12)); hold_idx.append(dates[d])
            if exec_at == "open":
                h *= intraday
            pending = None
        else:
            h *= gap * intraday

        V = h.sum() + cash
        eq[d] = V
        inv[d] = h.sum() / V if V > 0 else 0.0

        if sig_mask[d] and d + 1 < n_days:
            if TW is not None:
                w = np.nan_to_num(TW[d], nan=0.0)
                if np.isnan(TW[d]).all():
                    continue
            else:
                row = np.where(member[d], S[d], np.nan)
                held = h > 1e-12
                n_pick = top_n
                if dyn_n is not None:
                    n_pick = int(np.clip((row >= dyn_n[1]).sum(), dyn_n[0], top_n))
                idx = select_top(row, n_pick, held, buffer)
                w = np.zeros(n_assets)
                if len(idx) >= min(max(1, n_pick // 2), 20):
                    if weighting == "zvol":
                        w[idx] = zvol_weights(row, idx, vol[d, idx])
                    elif weighting == "invvol":
                        iv = 1.0 / np.where(vol[d, idx] > 0, vol[d, idx], np.nan)
                        iv = np.nan_to_num(iv, nan=np.nanmean(iv) if np.isfinite(iv).any() else 1.0)
                        w[idx] = iv / iv.sum()
                    elif weighting == "score":
                        s = np.clip(row[idx] - np.nanmin(row[idx]) + 1e-6, 0, None)
                        w[idx] = s / s.sum()
                    else:
                        w[idx] = 1.0 / len(idx)
            pending = w * float(np.clip(expo[d], 0.0, 1.0))

    equity = pd.Series(eq, index=dates)
    rets = equity.pct_change().fillna(0.0)
    holdings = pd.DataFrame(hold_rows, index=hold_idx, columns=panel.close.columns)
    return Result(equity, rets, pd.Series(turn, index=dates), holdings,
                  pd.Series(inv, index=dates), name)