"""
Exposure overlays: series in [0, 1] read on signal dates by engine.simulate
(``exposure=``). The un-invested part earns the panel's cash return.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from research.lab.engine import simulate


def ew_index(p) -> pd.Series:
    """Equal-weight index of PIT members (daily rebalanced)."""
    r = p.close.pct_change(fill_method=None).where(p.member.shift(1, fill_value=False))
    return (1 + r.mean(axis=1).fillna(0)).cumprod()


def trend(level: pd.Series, ma: int = 200, floor: float = 0.0) -> pd.Series:
    """1 when the index is above its moving average, else ``floor``."""
    above = level > level.rolling(ma, min_periods=ma).mean()
    return above.astype(float).where(above, floor)


def breadth(p, ma: int = 200, lo: float = 0.4, hi: float = 0.6) -> pd.Series:
    """Share of members above their own MA mapped linearly from lo→0 to hi→1."""
    c = p.close.ffill()
    above = (c > c.rolling(ma, min_periods=ma).mean()).where(p.member)
    share = above.mean(axis=1)
    return ((share - lo) / (hi - lo)).clip(0, 1)


def drawdown_brake(level: pd.Series, dd: float = 0.15, ma: int = 50) -> pd.Series:
    """Cut to zero when the index is > dd below its 1y high AND below its MA."""
    high = level.rolling(252, min_periods=60).max()
    bad = (level / high - 1 < -dd) & (level < level.rolling(ma, min_periods=ma).mean())
    return (~bad).astype(float)


def vol_target(p, score, target: float, lookback: int = 60, cap: float = 1.0,
               base_kwargs: dict | None = None, exposure: pd.Series | None = None) -> pd.Series:
    """Scale exposure by target / realized vol of the unscaled strategy.

    Two passes: run the strategy fully invested, measure its trailing
    daily-return vol, then size the next rebalance; any given ``exposure``
    (e.g. a trend filter) multiplies the result.
    """
    kw = dict(base_kwargs or {})
    base = simulate(p, score, **kw)
    rv = base.returns.rolling(lookback, min_periods=lookback // 2).std() * np.sqrt(252)
    scale = (target / rv).clip(upper=cap).fillna(cap)
    if exposure is not None:
        scale = scale * exposure.reindex(scale.index).ffill().fillna(1.0)
    return scale


# ---------------------------------------------------------------------------
# Literature overlays on an index level (all read at the close of the signal
# date; engine.simulate applies them at the next rebalance)
# ---------------------------------------------------------------------------

def faber_sma(level: pd.Series, months: int = 10) -> pd.Series:
    """Faber (2007): invested while the month-end close is above its
    ``months``-month simple average; the state changes only at month ends."""
    m = level.resample("ME").last()
    on = (m > m.rolling(months, min_periods=months).mean()).astype(float)
    on = on.where(m.rolling(months).count() >= months, 1.0)
    return on.reindex(level.index, method="ffill").fillna(1.0)


def derisk_vol(level: pd.Series, halflife: int = 20, min_hist: int = 252) -> pd.Series:
    """De-risk-only volatility target (Harvey et al. 2018 estimator): exposure
    = min(1, median vol to date / EWMA vol), so it only ever cuts."""
    r = level.pct_change(fill_method=None).fillna(0.0)
    v = np.sqrt((r ** 2).ewm(halflife=halflife, min_periods=halflife).mean()) * np.sqrt(252)
    med = v.expanding(min_periods=min_hist).median()
    return (med / v).clip(upper=1.0).fillna(1.0)


def momentum_state(level: pd.Series, months: int = 24, cut: float = 0.5) -> pd.Series:
    """Daniel & Moskowitz (2016) / Cooper et al. (2004): after a negative
    ``months``-month market return momentum is crash-prone; hold ``cut``."""
    m = level.resample("ME").last()
    down = (m / m.shift(months) - 1) < 0
    return down.map({True: cut, False: 1.0}).reindex(level.index, method="ffill").fillna(1.0)


# --- statistical jump model (Shu, Yu & Mulvey 2024, arXiv:2402.05272) -------

def jm_features(r: pd.Series) -> pd.DataFrame:
    """EWM downside deviation (halflife 10) and EWM Sortino ratios
    (halflives 20 and 60) of daily excess returns."""
    feats = {}
    dd10 = np.sqrt((np.minimum(r, 0) ** 2).ewm(halflife=10).mean())
    feats["dd10"] = np.log(dd10 + 1e-8)
    for hl in (20, 60):
        mu = r.ewm(halflife=hl).mean()
        dd = np.sqrt((np.minimum(r, 0) ** 2).ewm(halflife=hl).mean())
        feats[f"sortino{hl}"] = (mu / (dd + 1e-8)).clip(-10, 10)
    return pd.DataFrame(feats)


try:                                   # the DP loops are ~50x faster compiled
    from numba import njit
except ImportError:                    # pragma: no cover - pure-Python fallback
    def njit(f=None, **kw):
        return f if f is not None else (lambda g: g)


@njit(cache=True)
def _jm_viterbi(cost: np.ndarray, lam: float) -> np.ndarray:
    T = cost.shape[0]
    V0, V1 = cost[0, 0], cost[0, 1]
    back = np.zeros((T, 2), np.int64)
    for t in range(1, T):
        a0, b0 = V0, V1 + lam          # into state 0: stay / switch
        a1, b1 = V1, V0 + lam
        back[t, 0] = 0 if a0 <= b0 else 1
        back[t, 1] = 1 if a1 <= b1 else 0
        V0, V1 = cost[t, 0] + min(a0, b0), cost[t, 1] + min(a1, b1)
    s = np.empty(T, np.int64)
    s[T - 1] = 0 if V0 <= V1 else 1
    for t in range(T - 1, 0, -1):
        s[t - 1] = back[t, s[t]]
    return s


@njit(cache=True)
def _jm_forward(cost: np.ndarray, lam: float) -> np.ndarray:
    """State at each t from a forward DP over data through t (no look-ahead)."""
    T = cost.shape[0]
    V0, V1 = cost[0, 0], cost[0, 1]
    out = np.empty(T, np.int64)
    out[0] = 0 if V0 <= V1 else 1
    for t in range(1, T):
        V0, V1 = cost[t, 0] + min(V0, V1 + lam), cost[t, 1] + min(V1, V0 + lam)
        out[t] = 0 if V0 <= V1 else 1
    return out


def _jm_cost(X: np.ndarray, mu: np.ndarray) -> np.ndarray:
    return ((X[:, None, :] - mu[None]) ** 2).sum(axis=2)


def _jm_fit(X: np.ndarray, lam: float, n_init: int = 6, iters: int = 30, seed: int = 0):
    """Two-state jump model: min sum ||x_t - mu_{s_t}||^2 + lam * [s_t != s_{t-1}]
    by coordinate descent (dynamic programming for states, means for centers)."""
    rng = np.random.default_rng(seed)
    best = (np.inf, None)
    for _ in range(n_init):
        mu = X[rng.choice(len(X), 2, replace=False)]
        for _ in range(iters):
            s = _jm_viterbi(_jm_cost(X, mu), lam)
            new = np.array([X[s == k].mean(axis=0) if (s == k).any() else mu[k] for k in (0, 1)])
            if np.allclose(new, mu):
                break
            mu = new
        s = _jm_viterbi(_jm_cost(X, mu), lam)
        obj = _jm_cost(X, mu)[np.arange(len(X)), s].sum() + lam * np.abs(np.diff(s)).sum()
        if obj < best[0]:
            best = (obj, mu.copy())
    return best[1]


def _jm_online(X: np.ndarray, mu: np.ndarray, lam: float) -> np.ndarray:
    return _jm_forward(_jm_cost(X, mu), lam)


def jump_model(level: pd.Series, cash_ret: pd.Series | None = None,
               lams: tuple[float, ...] = (0.0, 5.0, 10.0, 20.0, 30.0, 50.0, 80.0, 120.0),
               train_days: int = 3000, min_train: int = 1500, refit_months: int = 6,
               select_years: int = 8, min_select_years: int = 4, cost: float = 0.001) -> pd.Series:
    """0/1 exposure from the jump model. Centers are refit every
    ``refit_months`` on the trailing ``train_days`` (at least ``min_train``);
    features are standardized with the training window's moments; the bull
    state is the one with the higher mean return in training. Each month the
    jump penalty is re-chosen as the one whose online 0/1 strategy (1-day
    delay, ``cost`` per switch) had the best Sharpe over the trailing
    ``select_years`` (at least ``min_select_years``)."""
    r = level.pct_change(fill_method=None).fillna(0.0)
    cr = cash_ret.reindex(r.index).fillna(0.0) if cash_ret is not None else pd.Series(0.0, index=r.index)
    ex = r - cr
    F = jm_features(ex)
    dates = r.index
    refits = [d for d in pd.date_range(dates[0], dates[-1], freq=f"{refit_months}MS")
              if dates.searchsorted(d) >= min_train]
    states = {}
    for lam in lams:
        st = pd.Series(np.nan, index=dates)
        for i, d in enumerate(refits):
            e = dates.searchsorted(d)
            s0 = max(0, e - train_days)
            Xtr = F.iloc[s0:e]
            m, sd = Xtr.mean(), Xtr.std().replace(0, 1)
            mu = _jm_fit(((Xtr - m) / sd).values, lam)
            nxt = dates.searchsorted(refits[i + 1]) if i + 1 < len(refits) else len(dates)
            # run the forward DP from the window start so the state is warm
            Xall = ((F.iloc[s0:nxt] - m) / sd).values
            path = _jm_online(Xall, mu, lam)
            tr_states = path[:e - s0]
            bull = int(np.argmax([ex.iloc[s0:e][tr_states == k].mean() if (tr_states == k).any() else -1
                                  for k in (0, 1)]))
            st.iloc[e:nxt] = (path[e - s0:] == bull).astype(float)
        states[lam] = st
    S = pd.DataFrame(states)
    # returns of each penalty's 0/1 strategy, held from the next day
    pos = S.shift(1)
    strat = pos.mul(r, axis=0) + (1 - pos).mul(cr, axis=0) - pos.diff().abs() * cost
    month_ends = strat.resample("ME").last().index
    choice = pd.Series(np.nan, index=dates)
    for me in month_ends:
        hist = strat.loc[:me].dropna()
        hist = hist[hist.index > me - pd.DateOffset(years=select_years)]
        if len(hist) < 252 * min_select_years:
            continue
        sh = (hist.mean() / hist.std()).dropna()
        if sh.empty:
            continue
        lam = sh.idxmax()
        nxt = dates.searchsorted(me, side="right")
        end = dates.searchsorted(me + pd.offsets.MonthEnd(1), side="right")
        choice.iloc[nxt:end] = S[lam].iloc[nxt:end].values
    return choice.fillna(1.0)