"""Pure, ex-ante building blocks for the strategy-level allocation engine (sub-project B2a).

This module hosts the pure allocation helpers plus `compute_allocation`, the ex-ante composition;
`target_capital` is added by a later B2a task. `full_history_vol` sizes from the FULL history
(stable/ex-ante, NOT a trailing window — the load-bearing anti-reactive property); `equal_weight_marginal_prune`
is the leave-one-out quality gate over the equal-weight book; `killswitch_active` reports the rare drawdown
killswitch's latch state; `compute_allocation` composes all three into ex-ante-sized, gated,
killswitch-aware target weights. No I/O; no reactive de-sizing."""
from __future__ import annotations

import math

import numpy as np

from fxhnt.application.allocation_policy import AllocationPolicy, _GAUSSIAN_ES
from fxhnt.application.promotion_corr import deploy_book_returns
from fxhnt.domain.gauntlet import per_period_sharpe

_ANN = np.sqrt(365.0)


def full_history_vol(returns: dict[str, float]) -> float:
    """Annualised (√365) std of return values over the WHOLE series.

    Returns 0.0 if fewer than 2 observations.
    This is ex-ante stable: uses the full history, NOT a trailing window.

    Annualises by the CALENDAR-DAY convention (√365), consistent with `target_vol`'s convention, so the
    RELATIVE leverage K is correct across the mixed crypto/ETF book — even though a weekday-only series
    (e.g. an ETF) is slightly over-annualised in absolute terms.
    """
    vals = list(returns.values())
    if len(vals) < 2:
        return 0.0
    return float(np.std(np.array(vals)) * _ANN)


def inverse_vol_weights(returns_by_strategy: dict[str, dict[str, float]]) -> dict[str, float]:
    """Risk-parity base: weight each strategy ∝ 1/full_history_vol, normalized to sum 1. Zero/degenerate-vol
    strategy → 0; all-degenerate → equal weight. Reuses the ex-ante full-history vol (never trailing)."""
    inv = {s: (1.0 / v if (v := full_history_vol(r)) > 0.0 else 0.0) for s, r in returns_by_strategy.items()}
    total = sum(inv.values())
    if total <= 0.0:
        n = len(returns_by_strategy)
        return {s: 1.0 / n for s in returns_by_strategy} if n else {}
    return {s: v / total for s, v in inv.items()}


def cvar_daily(returns: dict[str, float], alpha: float) -> float:
    """Expected shortfall (CVaR) of the daily return series at tail level `alpha`, over the WHOLE history
    (ex-ante, NOT trailing) — the mean of the worst ceil(alpha*n) returns, as a POSITIVE loss magnitude.

    0.0 if fewer than 2 observations OR the worst tail is non-negative (all-gains) — so the CVaR leverage
    cap is inactive rather than dividing by ~0. Uses the same full-history discipline as `full_history_vol`."""
    vals = sorted(returns.values())
    n = len(vals)
    if n < 2:
        return 0.0
    n_tail = max(1, math.ceil(alpha * n))
    tail_mean = float(np.mean(vals[:n_tail]))   # sorted ascending -> worst returns first
    return max(0.0, -tail_mean)


def _book_sharpe(returns_by_strategy: dict[str, dict[str, float]]) -> float:
    """Helper: compute the equal-weight book's per-period Sharpe ratio.

    Returns 0.0 if fewer than 2 observations.
    """
    book = deploy_book_returns(returns_by_strategy)  # equal-weight per-date average
    if len(book) < 2:
        return 0.0
    return float(per_period_sharpe(np.array([book[d] for d in sorted(book)])))


def equal_weight_marginal_prune(
    returns_by_strategy: dict[str, dict[str, float]],
) -> set[str]:
    """Leave-one-out: bench a strategy only when removing it raises the equal-weight book's Sharpe.

    Shrinks to ≥1 strategy.
    """
    kept = set(returns_by_strategy)
    while len(kept) > 1:
        base = _book_sharpe({e: returns_by_strategy[e] for e in kept})
        drop = next(
            (
                e
                for e in sorted(kept)
                if _book_sharpe({x: returns_by_strategy[x] for x in kept - {e}})
                > base + 1e-9
            ),
            None,
        )
        if drop is None:
            break
        kept.discard(drop)
    return kept


def killswitch_active(
    book_returns: dict[str, float], kill_dd: float, reenter_dd: float
) -> bool:
    """The final latch state after folding the drawdown killswitch over the sorted series.

    Activates when drawdown exceeds kill_dd; deactivates when drawdown falls below reenter_dd.
    """
    killed = False
    shadow_eq = 1.0
    peak = 1.0
    for d in sorted(book_returns):
        shadow_eq *= 1.0 + book_returns[d]
        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 killed


def compute_allocation(
    returns_by_strategy: dict[str, dict[str, float]], policy: AllocationPolicy,
    *, killswitch_returns: dict[str, dict[str, float]] | None = None,
) -> dict[str, float]:
    """The ex-ante composition: eligibility -> marginal gate -> risk-parity (inverse-vol, capped) -> killswitch -> vol-target K.

    K, the marginal gate, and the inverse-vol base derive from the FULL-history vol of the equal-weight
    survivor book over `returns_by_strategy` (the stable backtest sizing series; never a trailing window) —
    the load-bearing anti-reactive invariant.

    The drawdown killswitch is LIVE protection, so it runs on `killswitch_returns` (the forward/live record)
    when provided — NOT the sizing series. A backtest sizing series can carry a historical >kill_dd drawdown
    that must never permanently disable an edge going forward. When `killswitch_returns is None` the
    killswitch falls back to `returns_by_strategy` (byte-identical to before this argument existed). Survivors
    absent from the live record collapse to an empty book -> `killswitch_active` returns False (no kill).

    Returns {} if no strategy has enough history; all-zero weights if the rare aggregate killswitch is latched.
    """
    # 1. eligibility: a stable ex-ante vol needs enough history
    eligible = {s: r for s, r in returns_by_strategy.items() if len(r) >= policy.min_history_days}
    if not eligible:
        return {}
    # 2. marginal-Sharpe quality gate (equal-weight book), 3. equal-weight survivors
    survivors = equal_weight_marginal_prune(eligible) if policy.marginal_gate else set(eligible)
    # risk-parity base (inverse full-history vol), capped at the per-strategy concentration limit. NO
    # renormalise: when the cap binds the book deliberately holds cash rather than over-concentrating.
    rp = inverse_vol_weights({s: eligible[s] for s in survivors})
    base = {s: min(rp.get(s, 0.0), policy.max_strategy_weight) for s in survivors}
    # equal-weight survivor book (backtest sizing series) -> full-history vol for K below.
    book = deploy_book_returns({s: eligible[s] for s in survivors})
    # 5. rare latching aggregate killswitch -> flatten. Runs on the LIVE record (killswitch_returns) when
    #    injected, else the sizing series (byte-identical fallback). An empty ks_book -> not latched.
    ks_src = killswitch_returns if killswitch_returns is not None else returns_by_strategy
    ks_book = deploy_book_returns({s: ks_src[s] for s in survivors if s in ks_src})
    if killswitch_active(ks_book, policy.kill_dd, policy.reenter_dd):
        return {s: 0.0 for s in survivors}
    # 4. ex-ante vol-target leverage from the FULL-history book vol (never trailing), floored by a CVaR
    #    tail cap so fat-left-tail books de-lever. target_cvar is the Gaussian ES at target_vol, so the CVaR
    #    cap is a NO-OP for a Gaussian book and only bites fatter-than-normal left tails.
    if policy.cvar_alpha not in _GAUSSIAN_ES:   # fail-closed regardless of tail sign (not only when cv>0)
        raise ValueError(f"cvar_alpha {policy.cvar_alpha!r} has no Gaussian-ES multiplier; supported: {sorted(_GAUSSIAN_ES)}")
    vol = full_history_vol(book)
    k_vol = min(policy.max_leverage, policy.kelly_fraction * policy.target_vol / vol) if vol > 1e-12 else 0.0
    cv = cvar_daily(book, policy.cvar_alpha)
    if cv > 1e-12:
        target_cvar_daily = _GAUSSIAN_ES[policy.cvar_alpha] * policy.target_vol / math.sqrt(365.0)
        k_cvar = min(policy.max_leverage, policy.kelly_fraction * target_cvar_daily / cv)
    else:
        k_cvar = policy.max_leverage   # CVaR inactive -> don't constrain below k_vol
    k = min(k_vol, k_cvar)
    return {s: k * w for s, w in base.items()}


def target_capital(target_weights: dict[str, float], equity: float,
                   policy: AllocationPolicy, *,
                   capacity_ceiling: dict[str, float | None] | None = None) -> dict[str, float]:
    """Per-strategy signed dollars. Total deployed gross = min(equity·Σ|w|, capital_ceiling) — the HARD
    human gate: never deploy more than the ceiling. Signs preserved.

    `capacity_ceiling` (optional) is a per-strategy dollar ceiling from the honest-reference capacity curve:
    after the fund-level ceiling scaling, each strategy's MAGNITUDE is additionally capped at its own finite
    ceiling (sign preserved: a short caps at -min(|d|, ceiling)). A `None`/missing/non-finite ceiling for a
    strategy is skipped. `capacity_ceiling=None` (or empty) → byte-identical to the fund-ceiling-only result."""
    if not target_weights:
        return {}
    gross = sum(abs(w) for w in target_weights.values()) * equity
    if not math.isfinite(gross) or gross <= 0.0:
        return {s: 0.0 for s in target_weights}
    deployed = min(gross, policy.capital_ceiling)
    scale = deployed / gross
    dollars = {s: w * equity * scale for s, w in target_weights.items()}
    if capacity_ceiling:
        for s, d in dollars.items():
            ceil = capacity_ceiling.get(s)
            if ceil is not None and math.isfinite(ceil):
                dollars[s] = math.copysign(min(abs(d), ceil), d)
    return dollars


def apply_venue_caps(dollars: dict[str, float], venue_by_sid: dict[str, str],
                     max_venue_weight: float, equity: float) -> dict[str, float]:
    """Cap each venue's total deployed gross at `max_venue_weight * equity` (sign-preserving, proportional
    scale-down within the venue). `max_venue_weight >= 1.0` (or non-finite equity) -> no-op. A sid missing from
    `venue_by_sid` falls in the "" bucket (its own venue)."""
    if not dollars or max_venue_weight >= 1.0 or not math.isfinite(equity) or equity <= 0.0:
        return dollars
    cap = max_venue_weight * equity
    by_venue: dict[str, list[str]] = {}
    for sid in dollars:
        by_venue.setdefault(venue_by_sid.get(sid, ""), []).append(sid)
    for sids in by_venue.values():
        gross = sum(abs(dollars[s]) for s in sids)
        if gross > cap and gross > 0.0:
            scale = cap / gross
            for s in sids:
                dollars[s] *= scale
    return dollars


def capacity_ceiling_from_curve(curve: dict[float, float], min_sharpe: float) -> float | None:
    """Largest AUM whose honest Sharpe >= min_sharpe (the capacity ceiling).

    `None` ONLY for an EMPTY curve (no capacity information -> caller/`target_capital` treats it as
    unbounded). An edge that is unviable at EVERY tested AUM returns 0.0 — ZERO capacity, the OPPOSITE of
    unbounded — so `target_capital` caps its dollars to $0 (`min(|d|, 0) == 0`). Report the ceiling next to
    the curve so the threshold is visible."""
    if not curve:
        return None
    ok = [aum for aum in sorted(curve) if curve[aum] >= min_sharpe]
    return max(ok) if ok else 0.0
