"""Book-of-books allocator: combine date-stamped edge return series into one risk-managed
book via inverse-vol/risk-parity weights (per-sleeve cap) + fractional-Kelly leverage +
portfolio vol-target. Simple robust method (research-backed), NOT ML. No I/O."""
from __future__ import annotations

import math
from typing import Any

import numpy as np

from fxhnt.domain.backtest import compute_stats
from fxhnt.domain.gauntlet.core import evaluate, per_period_sharpe

# {epoch_day: daily_net_return} for one edge
EdgeReturns = dict[int, float]

_MAX_LEVERAGE = 3.0   # hard cap on the vol-target leverage (no full-Kelly blow-up)


def align_edges(series_by_edge: dict[str, EdgeReturns]) -> tuple[list[int], dict[str, np.ndarray]]:
    """Align every edge to the sorted UNION of all dates where AT LEAST one edge trades.

    Returns (dates, {edge: array}). Each edge's array is parallel to `dates`, using 0.0 where
    the edge has no return that day (a flat / uninvested day for that sleeve).
    """
    dates = sorted({d for s in series_by_edge.values() for d in s})
    idx = {d: i for i, d in enumerate(dates)}
    arrs: dict[str, np.ndarray] = {}
    for edge, series in series_by_edge.items():
        a = np.zeros(len(dates), dtype=float)
        for d, r in series.items():
            a[idx[d]] = r
        arrs[edge] = a
    return dates, arrs


def _inverse_vol_weights(window: dict[str, np.ndarray], max_sleeve_weight: float) -> dict[str, float]:
    """Inverse-vol (risk-parity-lite) weights from a trailing window, each capped then renormalized.

    weight_i ∝ 1/vol_i (normalized to sum 1), then each weight is capped at `max_sleeve_weight`
    and the book is renormalized. The cap stops a single low-vol/high-Sharpe sleeve dominating.
    """
    inv: dict[str, float] = {}
    for edge, arr in window.items():
        v = float(np.std(arr)) if arr.size > 1 else 0.0
        inv[edge] = (1.0 / v) if v > 1e-12 else 0.0
    total = sum(inv.values())
    if total <= 0.0:                              # no usable vol signal yet → equal weight
        n = len(window)
        return {e: 1.0 / n for e in window} if n else {}
    w: dict[str, float] = {e: inv[e] / total for e in window}
    # Iteratively cap + renormalize so that no sleeve exceeds max_sleeve_weight while the
    # uncapped sleeves still sum the book to 1 (standard cap-and-renormalize loop).
    for _ in range(len(window)):
        over = {e: x for e, x in w.items() if x > max_sleeve_weight + 1e-12}
        if not over:
            break
        capped_mass = max_sleeve_weight * len(over)
        free = {e: x for e, x in w.items() if e not in over}
        free_total = sum(free.values())
        remaining = max(0.0, 1.0 - capped_mass)
        new_w = {}
        for e in w:
            if e in over:
                new_w[e] = max_sleeve_weight
            elif free_total > 1e-12:
                new_w[e] = free[e] / free_total * remaining
            else:
                new_w[e] = 0.0
        w = new_w
    return w


def current_sleeve_weights(series_by_edge: dict[str, EdgeReturns], *, vol_lookback: int = 60,
                           max_sleeve_weight: float = 0.4) -> dict[str, float]:
    """Per-sleeve inverse-vol weights for the LATEST trailing window — the SAME computation
    `combine_book` does at a rebalance, exposed as a public accessor for the paper snapshot.

    Aligns the edges to the union date axis (0.0-fill on flat days), takes the last `vol_lookback`
    rows, and runs the existing `_inverse_vol_weights` (capped at `max_sleeve_weight`, renormalized).
    Does NOT change `combine_book`. Returns {} for an empty/no-date book.

    NOTE: this mirrors the UN-gated inverse-vol path — it does NOT apply `_marginal_prune` (the
    correlation-aware pruning `combine_book` can run), matching the current production call site which
    drives the paper snapshot off the raw inverse-vol weights.
    """
    dates, arrs = align_edges(series_by_edge)
    if not dates:
        return {}
    lo = max(0, len(dates) - vol_lookback)
    window = {e: a[lo:] for e, a in arrs.items()}
    return _inverse_vol_weights(window, max_sleeve_weight)


def _marginal_prune(window: dict[str, np.ndarray], max_sleeve_weight: float) -> dict[str, np.ndarray]:
    """Leave-one-out marginal-Sharpe gate: greedily bench a sleeve only when REMOVING it raises the
    inverse-vol book's trailing per-period Sharpe. Unlike a raw Sharpe floor this KEEPS a legitimate
    anticorrelated hedge (removing it lowers book Sharpe) while benching a dead non-diversifying sleeve
    (removing it raises it). Shrinks to at least one sleeve. Pure on the trailing window — no lookahead."""
    def book_sharpe(keys: set[str]) -> float:
        w = _inverse_vol_weights({e: window[e] for e in keys}, max_sleeve_weight)
        pw = np.zeros(len(next(iter(window.values()))), dtype=float)
        for e in keys:
            pw = pw + w[e] * window[e]
        return float(per_period_sharpe(pw))

    kept = set(window)
    while len(kept) > 1:
        base = book_sharpe(kept)
        drop = next((e for e in kept if book_sharpe(kept - {e}) > base + 1e-9), None)
        if drop is None:
            break
        kept.discard(drop)
    return {e: window[e] for e in window if e in kept}


def combine_book(series_by_edge: dict[str, EdgeReturns], *, vol_lookback: int = 60,
                 target_vol: float = 0.12, kelly_fraction: float = 0.5,
                 max_sleeve_weight: float = 0.4, periods_per_year: int = 365,
                 rebalance_every: int = 21,
                 min_sleeve_sharpe: float | None = None,
                 marginal_gate: bool = False) -> dict[int, float]:
    """Combine date-stamped edges into ONE risk-managed daily book return series.

    Method (uses ONLY trailing data — no lookahead):
      1. Align edges to the union date axis (0.0-fill on flat days).
      2. OPTIONAL EDGE-QUALITY GATE (`min_sleeve_sharpe`, default None = OFF): when set, every
         `rebalance_every` days bench any sleeve whose trailing per-period Sharpe over the window is
         below the floor, so a dead/decaying low-vol sleeve cannot earn inverse-vol weight and be
         levered up by the vol-target (the allocator's exploitable flaw — a dead cta-trend sleeve
         once inflated the book Sharpe 3.27→3.70 purely via manufactured leverage).
         CAVEAT — a naive Sharpe floor cannot distinguish a dead sleeve from a legitimate
         anticorrelated HEDGE (low standalone Sharpe but vol-reducing): it benches both. So it is
         OFF by default; a correct quality/decay layer (gate on a long window + correlation-benefit,
         or deflated-Sharpe) is future work. Use only for return-seeking sleeves you trust are edges.
         The CORRECT quality layer is `marginal_gate` (default False): leave-one-out — bench a sleeve
         only when REMOVING it raises the trailing book Sharpe. This keeps anticorrelated hedges and
         benches only dead, non-diversifying sleeves (it distinguishes the two; the Sharpe floor cannot).
         Then recompute inverse-vol weights over the ELIGIBLE sleeves only, CAP each at
         `max_sleeve_weight`, renormalize. Benched sleeves get weight 0. Hold between rebalances.
      3. Portfolio raw daily return r_p[t] = Σ_i w_i · r_i[t].
      4. Vol-target + fractional-Kelly leverage (recomputed each rebalance from trailing data):
             lev = kelly_fraction · (target_vol_daily / trailing_port_vol),  capped to [0, _MAX_LEVERAGE]
         where target_vol_daily = target_vol / sqrt(periods_per_year) and trailing_port_vol is the
         realized daily std of r_p over the trailing `vol_lookback` window. fractional-Kelly is the
         overall scalar so we run a fraction of the full vol-targeted exposure.
      5. book[t] = lev · r_p[t]. If NO sleeve is eligible, the book is flat (0.0) that period.

    Returns {epoch_day: book_return}.
    """
    dates, arrs = align_edges(series_by_edge)
    n = len(dates)
    if n == 0:
        return {}
    edges = list(arrs)
    # raw portfolio return per day under the currently-held weights
    target_vol_daily = target_vol / math.sqrt(periods_per_year)
    book: dict[int, float] = {}
    weights: dict[str, float] = {e: 1.0 / len(edges) for e in edges} if edges else {}
    leverage = 0.0
    for t in range(n):
        if t % rebalance_every == 0 and t > 1:     # need trailing data to estimate vol/edge
            lo = max(0, t - vol_lookback)
            # OPTIONAL edge-quality gate: keep only sleeves with usable vol and (if gating) a
            # trailing Sharpe >= floor. With the gate OFF (default) all sleeves with vol are eligible.
            eligible = {e: arrs[e][lo:t] for e in edges
                        if (arrs[e][lo:t].size > 1 and float(np.std(arrs[e][lo:t])) > 1e-12
                            and (min_sleeve_sharpe is None
                                 or per_period_sharpe(arrs[e][lo:t]) >= min_sleeve_sharpe))}
            if eligible and marginal_gate and len(eligible) >= 2:
                eligible = _marginal_prune(eligible, max_sleeve_weight)
            if eligible:
                ew = _inverse_vol_weights(eligible, max_sleeve_weight)
                weights = {e: ew.get(e, 0.0) for e in edges}
                # trailing raw portfolio vol under the NEW weights (no lookahead: uses [lo, t))
                port_window = np.zeros(t - lo, dtype=float)
                for e in edges:
                    port_window += weights[e] * arrs[e][lo:t]
                pv = float(np.std(port_window)) if port_window.size > 1 else 0.0
                leverage = min(_MAX_LEVERAGE, kelly_fraction * (target_vol_daily / pv)) if pv > 1e-12 else 0.0
            else:                                  # every sleeve bleeding → stand flat
                weights = {e: 0.0 for e in edges}
                leverage = 0.0
        raw = sum(weights[e] * arrs[e][t] for e in edges)
        book[dates[t]] = leverage * raw
    return book


def apply_drawdown_killswitch(book: dict[int, float], *, kill_dd: float = 0.15,
                              reenter_dd: float = 0.07) -> dict[int, float]:
    """Regime kill-switch: flatten the book (return 0) once its running drawdown exceeds `kill_dd`,
    re-entering only after the drawdown recovers back inside `reenter_dd` (hysteresis avoids whipsaw).

    The drawdown SIGNAL tracks the UN-throttled book equity (so recovery is observable while flat);
    the OUTPUT is the throttled series. No lookahead: each day's output uses the kill-state set by
    strictly prior days. This is the small-capital tail control — survive a crypto-deleveraging by
    standing down, not by buying an expensive hedge. Returns {epoch_day: throttled_return}.
    """
    killed = False
    shadow_eq = 1.0
    peak = 1.0
    out: dict[int, float] = {}
    for d in sorted(book):
        r = book[d]
        out[d] = 0.0 if killed else r           # decision uses kill-state from prior days only
        shadow_eq *= (1.0 + r)                   # un-throttled equity = regime signal
        peak = max(peak, shadow_eq)
        dd = (peak - shadow_eq) / peak if peak > 0 else 0.0
        if not killed and dd > kill_dd:
            killed = True
        elif killed and dd < reenter_dd:
            killed = False
    return out


def _json_floats(d: dict[str, Any]) -> dict[str, Any]:
    """Coerce model_dump scalars to JSON-native: numpy → python float, NaN/inf → 0.0."""
    out: dict[str, Any] = {}
    for k, v in d.items():
        if isinstance(v, (np.floating, float)):
            fv = float(v)
            out[k] = fv if np.isfinite(fv) else 0.0
        elif isinstance(v, (np.integer,)):
            out[k] = int(v)
        else:
            out[k] = v
    return out


def _correlation_matrix(series_by_edge: dict[str, EdgeReturns]) -> dict[str, dict[str, float]]:
    """Pairwise correlation between edges over aligned days where BOTH sleeves are non-zero
    (i.e. both actually traded). JSON-safe: NaN/inf → 0.0; self-correlation = 1.0."""
    _, arrs = align_edges(series_by_edge)
    edges = list(arrs)
    cm: dict[str, dict[str, float]] = {a: {} for a in edges}
    for a in edges:
        for b in edges:
            if a == b:
                cm[a][b] = 1.0
                continue
            mask = (arrs[a] != 0.0) & (arrs[b] != 0.0)
            xa, xb = arrs[a][mask], arrs[b][mask]
            if xa.size > 2 and xa.std() > 1e-12 and xb.std() > 1e-12:
                c = float(np.corrcoef(xa, xb)[0, 1])
                cm[a][b] = c if np.isfinite(c) else 0.0
            else:
                cm[a][b] = 0.0
    return cm


def evaluate_book(series_by_edge: dict[str, EdgeReturns], *, periods_per_year: int = 365,
                  oos_fraction: float = 0.40, dsr_min: float = 0.95,
                  oos_min_sharpe: float = 0.0, max_is_oos_decay: float = 0.50,
                  **combine_kwargs: Any) -> dict[str, Any]:
    """Combine the edges into one book, then gauntlet the combined series.

    Returns a JSON-safe report: combined `stats` (cagr/ann_vol/sharpe/max_drawdown/n_obs), the
    `verdict` (IS/OOS split, n_trials=1 — a single combined book), the pairwise `correlation_matrix`
    between edges, and each edge's `standalone_sharpe`. JSON-NaN-coerced like evaluate_funding_matrix.
    """
    book = combine_book(series_by_edge, periods_per_year=periods_per_year, **combine_kwargs)
    dates = sorted(book)
    r = np.array([book[d] for d in dates], dtype=float)
    stats = compute_stats(r, periods_per_year=periods_per_year)

    if len(r) < 3:
        verdict = {"passed": False, "dsr": 0.0, "is_sharpe": 0.0, "oos_sharpe": 0.0,
                   "n_trials": 1, "reasons": ["insufficient data"], "pvalue": 1.0}
    else:
        split = int((1.0 - oos_fraction) * len(r))
        is_r = r[:split]
        # n_trials=1: a single combined book, no multiple-testing search. sr_variance via a small
        # bootstrap of the IS per-period Sharpe (sampling-uncertainty proxy; falls back to 0.0).
        sr_variance = _is_sharpe_bootstrap_var(is_r)
        v = evaluate(is_r, r[split:], n_trials=1, sr_variance=sr_variance,
                     dsr_min=dsr_min, oos_min_sharpe=oos_min_sharpe,
                     max_is_oos_decay=max_is_oos_decay,
                     has_economic_rationale=True)  # combination of pre-vetted structural edges
        verdict = _json_floats(v.model_dump())

    standalone = {e: float(per_period_sharpe(np.array(list(s.values()), dtype=float)))
                  for e, s in series_by_edge.items()}
    standalone = {e: (x if np.isfinite(x) else 0.0) for e, x in standalone.items()}

    return {
        "stats": _json_floats(stats.model_dump()),
        "verdict": verdict,
        "correlation_matrix": _correlation_matrix(series_by_edge),
        "standalone_sharpe": standalone,
        "settings": {"periods_per_year": periods_per_year, "oos_fraction": oos_fraction,
                     **{k: v for k, v in combine_kwargs.items()}},
    }


def _is_sharpe_bootstrap_var(is_returns: np.ndarray, n_boot: int = 200) -> float:
    """Deterministic bootstrap variance of the IS per-period Sharpe (sampling uncertainty proxy)."""
    r = np.asarray(is_returns, float)
    r = r[np.isfinite(r)]
    if len(r) < 8:
        return 0.0
    rng = np.random.default_rng(0)               # deterministic
    srs = []
    for _ in range(n_boot):
        sample = rng.choice(r, size=len(r), replace=True)
        srs.append(per_period_sharpe(sample))
    v = float(np.var(srs))
    return v if np.isfinite(v) else 0.0
