"""
Markov-chain trading model (equity / ETF adaptation).

Adapts the Polymarket state-machine framework to tradable securities by mapping
**rolling price percentile** (0–1) into discrete states, estimating transition
probabilities from history, and simulating forward paths with Monte Carlo.

Core steps (mirrors the quant article):
  1. Build transition matrix from discretized price-percentile history
  2. Monte Carlo paths to horizon → raw probability of "bullish" resolution
  3. Optional calibration (Polymarket longshot table or equity shrinkage)
  4. Kelly sizing vs current percentile ("market price")
  5. Walk-forward: re-estimate matrix using only past data at each step
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import pandas as pd

# Becker (2026) Polymarket empirical calibration — optional for prediction markets.
POLYMARKET_CALIBRATION: dict[float, float] = {
    0.01: 0.0043,
    0.05: 0.0418,
    0.10: 0.087,
    0.20: 0.181,
    0.30: 0.285,
    0.50: 0.500,
    0.70: 0.715,
    0.80: 0.819,
    0.90: 0.913,
    0.95: 0.958,
}


def rolling_percentile_rank(close: pd.Series, window: int) -> pd.Series:
    """Percentile of today's close within the trailing ``window``-day range (0–1)."""
    c = close.astype(np.float64)
    lo = c.rolling(window, min_periods=window).min()
    hi = c.rolling(window, min_periods=window).max()
    span = (hi - lo).replace(0.0, np.nan)
    return ((c - lo) / span).clip(0.0, 1.0)


def discretize_states(values: np.ndarray, n_states: int) -> np.ndarray:
    """Map continuous values in [0, 1] to integer states ``0 .. n_states-1``."""
    arr = np.asarray(values, dtype=np.float64)
    states = np.full(len(arr), n_states // 2, dtype=int)
    valid = np.isfinite(arr)
    if valid.any():
        states[valid] = np.clip((arr[valid] * n_states).astype(int), 0, n_states - 1)
    return states


def build_transition_matrix(states: np.ndarray, n_states: int) -> np.ndarray:
    """
    Count state→state transitions and row-normalize to probabilities.

    Rows with zero observations fall back to uniform distribution.
    """
    T = np.zeros((n_states, n_states), dtype=np.float64)
    s = np.asarray(states, dtype=int)
    for i in range(len(s) - 1):
        T[s[i], s[i + 1]] += 1.0
    row_sums = T.sum(axis=1, keepdims=True)
    empty = row_sums.squeeze() == 0
    if empty.any():
        T[empty, :] = 1.0
        row_sums = T.sum(axis=1, keepdims=True)
    return T / row_sums


def smooth_transition_matrix(T: np.ndarray, min_count: float = 20.0) -> np.ndarray:
    """
    Laplace-style smoothing for sparse rows (< ``min_count`` transitions).

    Blends each sparse row toward uniform to avoid noise-dominated estimates.
    """
    n = len(T)
    counts = T * T.sum(axis=1, keepdims=True)  # approximate if already normalized
    row_totals = counts.sum(axis=1)
    out = T.copy()
    uniform = np.full(n, 1.0 / n)
    for i in range(n):
        if row_totals[i] < min_count:
            alpha = min_count / max(row_totals[i], 1.0)
            out[i] = (out[i] + alpha * uniform) / (1.0 + alpha)
    out = out / out.sum(axis=1, keepdims=True)
    return out


def monte_carlo_probability(
    T: np.ndarray,
    start_state: int,
    *,
    horizon: int,
    n_sims: int = 10_000,
    rng: np.random.Generator | None = None,
    bullish_threshold_state: int | None = None,
) -> float:
    """
    Fraction of simulated paths ending in "bullish" territory (state >= midpoint).

    Vectorized over ``n_sims`` paths for speed on walk-forward backtests.
    """
    n_states = len(T)
    if bullish_threshold_state is None:
        bullish_threshold_state = n_states // 2
    gen = rng if rng is not None else np.random.default_rng()

    cdf = np.cumsum(T, axis=1)
    cdf[:, -1] = 1.0
    states = np.full(n_sims, int(start_state), dtype=np.int64)

    for _ in range(horizon):
        u = gen.random(n_sims)
        nxt = np.empty(n_sims, dtype=np.int64)
        for s in range(n_states):
            mask = states == s
            if mask.any():
                nxt[mask] = np.searchsorted(cdf[s], u[mask], side="right")
        states = nxt

    return float((states >= bullish_threshold_state).mean())


def calibrate_probability(raw_prob: float, table: dict[float, float] | None = None) -> float:
    """Piecewise-linear calibration (Polymarket table by default)."""
    cal = table if table is not None else POLYMARKET_CALIBRATION
    keys = sorted(cal.keys())
    if raw_prob <= keys[0]:
        return cal[keys[0]]
    if raw_prob >= keys[-1]:
        return cal[keys[-1]]
    for i in range(len(keys) - 1):
        lo, hi = keys[i], keys[i + 1]
        if lo <= raw_prob <= hi:
            frac = (raw_prob - lo) / (hi - lo)
            return cal[lo] + frac * (cal[hi] - cal[lo])
    return raw_prob


def equity_shrink_calibration(raw_prob: float, shrink: float = 0.15) -> float:
    """Pull raw MC estimate toward 0.5 — mild anti-overconfidence for equities."""
    return (1.0 - shrink) * raw_prob + shrink * 0.5


def kelly_fraction_yes(p_win: float, market_price: float) -> float:
    """
    Full Kelly fraction for a binary contract priced at ``market_price`` (0–1).

    For equities we use ``market_price`` = current percentile rank.
    """
    cost = float(market_price)
    if cost <= 0.0 or cost >= 1.0:
        return 0.0
    b = (1.0 - cost) / cost
    p = float(p_win)
    q = 1.0 - p
    f = (b * p - q) / b
    return max(0.0, f)


def kelly_exposure(
    p_win: float,
    market_price: float,
    *,
    kelly_mult: float = 0.25,
    max_exposure: float = 1.0,
) -> float:
    """Quarter-Kelly (default) exposure capped at ``max_exposure``."""
    f = kelly_fraction_yes(p_win, market_price) * kelly_mult
    return float(min(max_exposure, max(0.0, f)))


def kelly_short_exposure(
    cal_prob: float,
    market_price: float,
    *,
    kelly_mult: float = 0.25,
    max_exposure: float = 1.0,
) -> float:
    """
    Kelly-sized SHORT exposure (returned as positive magnitude; caller negates).

    Mirrors ``kelly_exposure`` but from the perspective of a bear position:
    probability of falling = 1 - cal_prob, cost = 1 - market_price.
    """
    p_down = 1.0 - float(cal_prob)
    mp_short = 1.0 - float(market_price)
    f = kelly_fraction_yes(p_down, mp_short) * kelly_mult
    return float(min(max_exposure, max(0.0, f)))


@dataclass
class MarkovTradeSignal:
    """Output of a single-day Markov analysis for one ticker."""

    raw_prob: float
    calibrated_prob: float
    market_price: float
    edge: float
    exposure: float
    start_state: int
    decision: str  # LONG, CASH, PASS


@dataclass
class MarkovChainTradingModel:
    """
    Markov state machine for percentile-ranked price series.

    Parameters
    ----------
    n_states
        Number of discrete percentile buckets (default 10 → deciles).
    horizon
        Monte Carlo steps (trading days) until "resolution".
    n_sims
        Number of Monte Carlo paths.
    min_transitions
        Minimum transitions per row before smoothing kicks in.
    kelly_mult
        Fraction of full Kelly (0.25 = quarter-Kelly).
    edge_threshold
        Minimum |edge| to act (in probability units, e.g. 0.03 = 3¢).
    calibration_mode
        ``none`` | ``polymarket`` | ``equity_shrink``
    rng_seed
        Seed for reproducible Monte Carlo.
    """

    n_states: int = 10
    horizon: int = 21
    n_sims: int = 10_000
    min_transitions: float = 20.0
    kelly_mult: float = 0.25
    edge_threshold: float = 0.03
    calibration_mode: str = "equity_shrink"
    equity_shrink: float = 0.15
    rng_seed: int | None = 42
    _rng: np.random.Generator = field(init=False, repr=False)

    def __post_init__(self) -> None:
        self._rng = np.random.default_rng(self.rng_seed)

    def _calibrate(self, raw: float) -> float:
        if self.calibration_mode == "polymarket":
            return calibrate_probability(raw)
        if self.calibration_mode == "equity_shrink":
            return equity_shrink_calibration(raw, self.equity_shrink)
        return raw

    def analyze(
        self,
        percentile_history: np.ndarray,
        current_percentile: float,
    ) -> MarkovTradeSignal:
        """
        Run full pipeline on a history of percentile ranks ending yesterday.

        ``percentile_history`` should NOT include today's value (no lookahead).
        """
        hist = np.asarray(percentile_history, dtype=np.float64)
        hist = hist[~np.isnan(hist)]
        if len(hist) < self.n_states + 2:
            return MarkovTradeSignal(
                raw_prob=float("nan"),
                calibrated_prob=float("nan"),
                market_price=float(current_percentile),
                edge=float("nan"),
                exposure=0.0,
                start_state=self.n_states // 2,
                decision="PASS",
            )

        states = discretize_states(hist, self.n_states)
        T = build_transition_matrix(states, self.n_states)
        T = smooth_transition_matrix(T, min_count=self.min_transitions)

        start_state = int(
            discretize_states(np.array([current_percentile]), self.n_states)[0]
        )
        raw_prob = monte_carlo_probability(
            T,
            start_state,
            horizon=self.horizon,
            n_sims=self.n_sims,
            rng=self._rng,
        )
        cal_prob = self._calibrate(raw_prob)
        market_price = float(np.clip(current_percentile, 0.0, 1.0))
        edge = cal_prob - market_price

        if abs(edge) < self.edge_threshold or not np.isfinite(edge):
            return MarkovTradeSignal(
                raw_prob=raw_prob,
                calibrated_prob=cal_prob,
                market_price=market_price,
                edge=edge,
                exposure=0.0,
                start_state=start_state,
                decision="PASS",
            )

        if edge > 0:
            p_for_kelly = cal_prob
            exposure = kelly_exposure(
                p_for_kelly, market_price, kelly_mult=self.kelly_mult
            )
            # Scale exposure by edge strength; require minimum tilt to go long.
            exposure = min(1.0, exposure + edge)
            decision = "LONG" if exposure > 0.05 else "PASS"
        else:
            # Bearish edge → stay in cash for this sleeve.
            exposure = 0.0
            decision = "CASH"

        return MarkovTradeSignal(
            raw_prob=raw_prob,
            calibrated_prob=cal_prob,
            market_price=market_price,
            edge=edge,
            exposure=exposure if decision == "LONG" else 0.0,
            start_state=start_state,
            decision=decision,
        )


def build_regime_transition_matrices(
    states: np.ndarray,
    vix: np.ndarray,
    n_states: int,
    *,
    vix_lo: float = 15.0,
    vix_hi: float = 25.0,
) -> dict[str, np.ndarray]:
    """
    Build pooled and VIX-conditioned transition matrices.

    Transitions are tagged by VIX at the *from* day:
      - ``low``  : VIX < vix_lo
      - ``mid``  : vix_lo <= VIX <= vix_hi
      - ``high`` : VIX > vix_hi
      - ``pooled`` : all transitions (baseline)
    """
    s = np.asarray(states, dtype=int)
    v = np.asarray(vix, dtype=np.float64)
    keys = ("low", "mid", "high", "pooled")
    counts = {k: np.zeros((n_states, n_states), dtype=np.float64) for k in keys}

    for i in range(len(s) - 1):
        if not np.isfinite(v[i]):
            continue
        fr, to = s[i], s[i + 1]
        counts["pooled"][fr, to] += 1.0
        if v[i] < vix_lo:
            counts["low"][fr, to] += 1.0
        elif v[i] > vix_hi:
            counts["high"][fr, to] += 1.0
        else:
            counts["mid"][fr, to] += 1.0

    out: dict[str, np.ndarray] = {}
    totals: dict[str, float] = {}
    for k, c in counts.items():
        totals[k] = float(c.sum())
        row_sums = c.sum(axis=1, keepdims=True)
        empty = row_sums.squeeze() == 0
        T = c.copy()
        if empty.any():
            T[empty, :] = 1.0
            row_sums = T.sum(axis=1, keepdims=True)
        out[k] = T / row_sums
    out["_transition_counts"] = totals  # type: ignore[assignment]
    return out


def select_vix_regime_matrix(
    matrices: dict[str, np.ndarray],
    vix_level: float,
    *,
    vix_lo: float = 15.0,
    vix_hi: float = 25.0,
    min_regime_transitions: float = 15.0,
) -> tuple[np.ndarray, str]:
    """
    Pick transition matrix for current VIX with fallback if regime data is sparse.

    Returns (matrix, regime_label).
    """
    counts = matrices.get("_transition_counts", {})
    pooled = matrices["pooled"]

    if not np.isfinite(vix_level):
        return pooled, "pooled"

    if vix_level > vix_hi:
        label = "high"
    elif vix_level < vix_lo:
        label = "low"
    else:
        label = "mid"

    if float(counts.get(label, 0.0)) < min_regime_transitions:
        return pooled, "pooled"
    return matrices[label], label


def markov_exposure_decision(
    cal_prob: float,
    market_price: float,
    edge: float,
    *,
    edge_threshold: float,
    kelly_mult: float,
    allow_short: bool,
    max_short_exposure: float,
) -> tuple[float, str]:
    """Map Markov probabilities to signed exposure and decision label."""
    if abs(edge) < edge_threshold or not np.isfinite(edge):
        return 0.0, "PASS"
    if edge > 0:
        exposure = kelly_exposure(cal_prob, market_price, kelly_mult=kelly_mult)
        exposure = min(1.0, exposure + edge)
        decision = "LONG" if exposure > 0.05 else "PASS"
        if decision != "LONG":
            exposure = 0.0
        return exposure, decision
    if allow_short:
        short_size = kelly_short_exposure(
            cal_prob,
            market_price,
            kelly_mult=kelly_mult,
            max_exposure=max_short_exposure,
        )
        short_size = min(max_short_exposure, short_size + abs(edge))
        if short_size > 0.05:
            return -short_size, "SHORT"
        return 0.0, "PASS"
    return 0.0, "CASH"


def walk_forward_markov_features(
    close: pd.Series,
    *,
    percentile_window: int = 252,
    lookback: int = 252,
    model: MarkovChainTradingModel | None = None,
    matrix_refresh: int = 5,
    vix: pd.Series | None = None,
    vix_lo: float = 15.0,
    vix_hi: float = 25.0,
) -> pd.DataFrame:
    """
    Walk-forward Markov features without trading-rule filtering.

    Returns percentile, raw/calibrated probability, edge, and VIX regime label.
    Use ``markov_exposure_decision`` to apply thresholds and Kelly sizing offline.
    """
    m = model if model is not None else MarkovChainTradingModel()
    pct = rolling_percentile_rank(close, percentile_window)
    states_all = discretize_states(pct.values, m.n_states)
    idx = close.index
    rows: list[dict] = []

    vix_aligned: np.ndarray | None = None
    if vix is not None:
        vix_aligned = vix.reindex(idx).ffill().astype(np.float64).values

    warmup = max(percentile_window, lookback) + m.horizon + 5
    T_cached: np.ndarray | None = None
    regime_mats_cached: dict[str, np.ndarray] | None = None
    refresh = max(1, int(matrix_refresh))

    for i in range(warmup, len(idx)):
        dt = idx[i]
        cur = float(pct.iloc[i - 1])
        if np.isnan(cur):
            rows.append(
                {
                    "date": dt,
                    "percentile": np.nan,
                    "edge": np.nan,
                    "raw_prob": np.nan,
                    "calibrated_prob": np.nan,
                    "vix_regime": "",
                }
            )
            continue

        recompute = T_cached is None or ((i - warmup) % refresh == 0)
        if recompute:
            hist_start = max(0, i - lookback)
            pct_slice = pct.values[hist_start:i]
            valid_mask = np.isfinite(pct_slice)
            hist_states = states_all[hist_start:i][valid_mask]
            if len(hist_states) < m.n_states + 2:
                rows.append(
                    {
                        "date": dt,
                        "percentile": cur,
                        "edge": np.nan,
                        "raw_prob": np.nan,
                        "calibrated_prob": np.nan,
                        "vix_regime": "",
                    }
                )
                continue
            if vix_aligned is not None:
                vix_slice = vix_aligned[hist_start:i][valid_mask]
                regime_mats_cached = build_regime_transition_matrices(
                    hist_states,
                    vix_slice,
                    m.n_states,
                    vix_lo=vix_lo,
                    vix_hi=vix_hi,
                )
                for key in regime_mats_cached:
                    if key.startswith("_"):
                        continue
                    regime_mats_cached[key] = smooth_transition_matrix(
                        regime_mats_cached[key], min_count=m.min_transitions
                    )
                T_cached = regime_mats_cached["pooled"]
            else:
                regime_mats_cached = None
                T_cached = build_transition_matrix(hist_states, m.n_states)
                T_cached = smooth_transition_matrix(T_cached, min_count=m.min_transitions)

        vix_regime_label = "pooled"
        T_use = T_cached
        if regime_mats_cached is not None and vix_aligned is not None:
            vix_cur = float(vix_aligned[i - 1])
            T_use, vix_regime_label = select_vix_regime_matrix(
                regime_mats_cached,
                vix_cur,
                vix_lo=vix_lo,
                vix_hi=vix_hi,
                min_regime_transitions=m.min_transitions,
            )

        start_state = int(discretize_states(np.array([cur]), m.n_states)[0])
        raw_prob = monte_carlo_probability(
            T_use,
            start_state,
            horizon=m.horizon,
            n_sims=m.n_sims,
            rng=m._rng,
        )
        cal_prob = m._calibrate(raw_prob)
        market_price = float(np.clip(cur, 0.0, 1.0))
        edge = cal_prob - market_price

        rows.append(
            {
                "date": dt,
                "percentile": cur,
                "edge": edge,
                "raw_prob": raw_prob,
                "calibrated_prob": cal_prob,
                "vix_regime": vix_regime_label,
            }
        )

    if not rows:
        return pd.DataFrame()
    out = pd.DataFrame(rows)
    out["date"] = pd.to_datetime(out["date"])
    return out.set_index("date")


def walk_forward_exposure_series(
    close: pd.Series,
    *,
    percentile_window: int = 252,
    lookback: int = 252,
    model: MarkovChainTradingModel | None = None,
    matrix_refresh: int = 5,
    vix: pd.Series | None = None,
    vix_lo: float = 15.0,
    vix_hi: float = 25.0,
    allow_short: bool = False,
    max_short_exposure: float = 0.5,
) -> pd.DataFrame:
    """
    Walk-forward Markov exposure (0–1) for one ticker.

    ``matrix_refresh``: re-estimate transition matrix every N trading days
    (intermediate days reuse the last matrix — speeds backtests ~Nx).

    When ``vix`` is provided, builds separate transition matrices for
    VIX < ``vix_lo``, VIX > ``vix_hi``, and mid/pooled regimes.

    Returns DataFrame with columns: percentile, exposure, decision, edge, raw_prob.
    """
    m = model if model is not None else MarkovChainTradingModel()
    pct = rolling_percentile_rank(close, percentile_window)
    states_all = discretize_states(pct.values, m.n_states)
    idx = close.index
    rows: list[dict] = []

    vix_aligned: np.ndarray | None = None
    if vix is not None:
        vx = vix.reindex(idx).ffill().astype(np.float64).values
        vix_aligned = vx

    warmup = max(percentile_window, lookback) + m.horizon + 5
    T_cached: np.ndarray | None = None
    regime_mats_cached: dict[str, np.ndarray] | None = None
    refresh = max(1, int(matrix_refresh))

    for i in range(warmup, len(idx)):
        dt = idx[i]
        cur = float(pct.iloc[i - 1])
        if np.isnan(cur):
            rows.append(
                {
                    "date": dt,
                    "percentile": np.nan,
                    "exposure": 0.0,
                    "decision": "PASS",
                    "edge": np.nan,
                    "raw_prob": np.nan,
                    "calibrated_prob": np.nan,
                    "vix_regime": "",
                }
            )
            continue

        recompute = T_cached is None or ((i - warmup) % refresh == 0)
        if recompute:
            hist_start = max(0, i - lookback)
            pct_slice = pct.values[hist_start:i]
            valid_mask = np.isfinite(pct_slice)
            hist_states = states_all[hist_start:i][valid_mask]
            if len(hist_states) < m.n_states + 2:
                rows.append(
                    {
                        "date": dt,
                        "percentile": cur,
                        "exposure": 0.0,
                        "decision": "PASS",
                        "edge": np.nan,
                        "raw_prob": np.nan,
                        "calibrated_prob": np.nan,
                        "vix_regime": "",
                    }
                )
                continue
            if vix_aligned is not None:
                vix_slice = vix_aligned[hist_start:i][valid_mask]
                regime_mats_cached = build_regime_transition_matrices(
                    hist_states,
                    vix_slice,
                    m.n_states,
                    vix_lo=vix_lo,
                    vix_hi=vix_hi,
                )
                for key in regime_mats_cached:
                    if key.startswith("_"):
                        continue
                    regime_mats_cached[key] = smooth_transition_matrix(
                        regime_mats_cached[key], min_count=m.min_transitions
                    )
                T_cached = regime_mats_cached["pooled"]
            else:
                regime_mats_cached = None
                T_cached = build_transition_matrix(hist_states, m.n_states)
                T_cached = smooth_transition_matrix(T_cached, min_count=m.min_transitions)

        vix_regime_label = "pooled"
        T_use = T_cached
        if regime_mats_cached is not None and vix_aligned is not None:
            vix_cur = float(vix_aligned[i - 1])
            T_use, vix_regime_label = select_vix_regime_matrix(
                regime_mats_cached,
                vix_cur,
                vix_lo=vix_lo,
                vix_hi=vix_hi,
                min_regime_transitions=m.min_transitions,
            )

        start_state = int(discretize_states(np.array([cur]), m.n_states)[0])
        raw_prob = monte_carlo_probability(
            T_use,
            start_state,
            horizon=m.horizon,
            n_sims=m.n_sims,
            rng=m._rng,
        )
        cal_prob = m._calibrate(raw_prob)
        market_price = float(np.clip(cur, 0.0, 1.0))
        edge = cal_prob - market_price

        exposure, decision = markov_exposure_decision(
            cal_prob,
            market_price,
            edge,
            edge_threshold=m.edge_threshold,
            kelly_mult=m.kelly_mult,
            allow_short=allow_short,
            max_short_exposure=max_short_exposure,
        )

        rows.append(
            {
                "date": dt,
                "percentile": cur,
                "exposure": exposure,
                "decision": decision,
                "edge": edge,
                "raw_prob": raw_prob,
                "calibrated_prob": cal_prob,
                "vix_regime": vix_regime_label,
            }
        )

    if not rows:
        return pd.DataFrame()
    out = pd.DataFrame(rows)
    out["date"] = pd.to_datetime(out["date"])
    return out.set_index("date")