"""
Modulo de sizing de Leonex.

Implementa los dos pilares de gestion de riesgo de Thorp y Carver:

    1) KELLY CRITERION FRACCIONAL (Thorp, "Beat the Dealer", 1962)
       Para una apuesta con prob_win p y payoff b = TP/SL:
            f* = (p * b - q) / b     con q = 1 - p
       f* es el porcentaje del bankroll que maximiza el crecimiento
       logaritmico esperado. Es OPTIMO en el limite, pero muy sensible
       a errores de estimacion en p. Por eso usamos Half-Kelly (f*/2),
       que sacrifica ~25% del crecimiento esperado pero corta a la mitad
       el drawdown maximo esperado.

    2) VOLATILITY TARGETING (Carver, "Systematic Trading", 2015)
       El portfolio debe operar con una volatilidad anualizada objetivo
       (tipico: 20%). Si la vol implicita del portfolio (dadas las
       posiciones propuestas) excede el target, multiplicamos TODOS los
       sizes por un escalar < 1 para volver al objetivo. NUNCA apalancamos
       por encima de Kelly — vol_target solo RECORTA, nunca infla.

Politica de uso (orden de precedencia):
       Plan -> kelly_fraction()     define cuanto invertir por trade
            -> portfolio_vol()      estima vol agregada con correlacion
            -> vol_target_scalar()  recorta si excede el target
            -> [HRP en executor]    redistribuye entre activos correlados

Diseno: funciones puras, sin side-effects, faciles de testar.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import numpy as np


# ─────────────────────────────────────────────────────────────────────────────
# Constantes calibradas (Thorp / Carver / Lopez de Prado)
# ─────────────────────────────────────────────────────────────────────────────
DEFAULT_KELLY_FRACTION = 0.5         # Half-Kelly: recomendacion clasica de Thorp
DEFAULT_MAX_POSITION_PCT = 0.25      # nunca >25% del equity en un solo trade
DEFAULT_TARGET_VOL_ANNUAL = 0.20     # 20% vol anual portfolio (Carver default)
DEFAULT_AVG_CORRELATION = 0.30       # correlacion media entre activos si no la conocemos
DEFAULT_TRADING_DAYS_PER_YEAR = 252  # para anualizar sigma diaria
DEFAULT_TP_SL_RATIO = 2.0            # TP=2*sigma, SL=1*sigma -> payoff b = 2.0
DEFAULT_FALLBACK_SIGMA_DAILY = 0.02  # 2% diario si no tenemos sigma estimada


@dataclass
class SizingResult:
    """Resultado del sizing por ticker. Todos los pct son fracciones del equity."""
    ticker: str
    prob_win: float
    payoff_b: float
    sigma_daily: float
    sigma_annual: float
    kelly_full: float            # f* completo (sin aplicar half-kelly)
    kelly_fractional: float      # f* * kelly_fraction
    kelly_capped: float          # min(kelly_fractional, max_position_pct)
    portfolio_vol_estimate: float
    vol_scalar: float            # <=1: factor de reduccion por vol target
    final_pct: float             # peso final = kelly_capped * vol_scalar
    final_usd: float             # final_pct * equity
    notes: str = ""


# ─────────────────────────────────────────────────────────────────────────────
# Kelly Criterion
# ─────────────────────────────────────────────────────────────────────────────

def kelly_fraction(
    prob_win: float,
    payoff_b: float = DEFAULT_TP_SL_RATIO,
    kelly_fraction_mult: float = DEFAULT_KELLY_FRACTION,
    cap: float = DEFAULT_MAX_POSITION_PCT,
) -> tuple[float, float, float]:
    """Devuelve (kelly_full, kelly_fractional, kelly_capped).

    Args:
        prob_win: probabilidad de ganar (e.g. prob_win del Meta-Labeling)
        payoff_b: ratio entre ganancia esperada y perdida esperada.
                  Para Triple Barrier con TP=2sigma, SL=1sigma -> b=2.0
        kelly_fraction_mult: multiplicador (0.5 = Half-Kelly clasico).
        cap: tope absoluto al fractional Kelly (e.g. 0.25 = max 25% equity).

    Math:
        f* = (p*b - q) / b   con q = 1 - p
        Negativo -> apuesta perdedora (edge negativo) -> 0.

    Returns:
        kelly_full       =  f*
        kelly_fractional =  f* * kelly_fraction_mult
        kelly_capped     =  min(kelly_fractional, cap), nunca negativo.
    """
    if not (0.0 <= prob_win <= 1.0):
        return 0.0, 0.0, 0.0
    if payoff_b <= 0:
        return 0.0, 0.0, 0.0
    q = 1.0 - prob_win
    f_full = (prob_win * payoff_b - q) / payoff_b
    if f_full <= 0 or not math.isfinite(f_full):
        return 0.0, 0.0, 0.0
    f_frac = f_full * kelly_fraction_mult
    f_cap = min(max(f_frac, 0.0), max(cap, 0.0))
    return float(f_full), float(f_frac), float(f_cap)


# ─────────────────────────────────────────────────────────────────────────────
# Volatility
# ─────────────────────────────────────────────────────────────────────────────

def annualized_vol(
    sigma_daily: float,
    trading_days_per_year: int = DEFAULT_TRADING_DAYS_PER_YEAR,
) -> float:
    """sigma_anual = sigma_daily * sqrt(trading_days)."""
    if sigma_daily <= 0 or not math.isfinite(sigma_daily):
        return 0.0
    return float(sigma_daily * math.sqrt(trading_days_per_year))


def portfolio_vol_estimate(
    weights: list[float],
    sigmas_annual: list[float],
    correlations: Optional[np.ndarray] = None,
    avg_correlation: float = DEFAULT_AVG_CORRELATION,
) -> float:
    """Volatilidad anualizada del portfolio implicita.

    Si tenemos la matriz de correlacion completa, calculamos
        sigma_p = sqrt(w' Σ w)    con Σ = D R D, D=diag(sigmas), R=correlation
    Si no, asumimos correlacion media constante (default 0.3):
        var = Σ w_i^2 sigma_i^2 + 2 Σ_{i<j} w_i w_j sigma_i sigma_j avg_corr

    Esto es la aproximacion de Carver para pre-trade vol estimation.
    """
    if not weights or not sigmas_annual or len(weights) != len(sigmas_annual):
        return 0.0
    w = np.asarray(weights, dtype=float)
    s = np.asarray(sigmas_annual, dtype=float)
    if not w.size:
        return 0.0
    if correlations is not None:
        corr = np.asarray(correlations, dtype=float)
        if corr.shape != (w.size, w.size):
            correlations = None
    if correlations is None:
        # Aproximacion con correlacion constante
        n = w.size
        var_diag = float(np.sum(w * w * s * s))
        # cross terms: 2 * sum_{i<j} w_i w_j s_i s_j avg_corr
        cross = 0.0
        for i in range(n):
            for j in range(i + 1, n):
                cross += w[i] * w[j] * s[i] * s[j]
        var_cross = 2.0 * cross * float(avg_correlation)
        var = max(var_diag + var_cross, 0.0)
    else:
        D = np.diag(s)
        cov = D @ np.asarray(correlations, dtype=float) @ D
        var = float(w @ cov @ w)
        var = max(var, 0.0)
    return float(math.sqrt(var))


def vol_target_scalar(
    current_portfolio_vol: float,
    target_vol: float = DEFAULT_TARGET_VOL_ANNUAL,
) -> float:
    """Factor multiplicativo de reduccion.

    - Si vol actual <= target: devolvemos 1.0 (no apalancamos por encima de Kelly).
    - Si vol actual >  target: devolvemos target / actual (< 1, recorta).
    - Si vol actual <= 0: devolvemos 0.0 (defensivo).

    NOTA: NUNCA superamos 1.0. Carver llama a esto "skipping leverage" — si la
    vol implicita ya es < target, no vamos a forzar leverage para llenarla.
    """
    if current_portfolio_vol <= 0 or not math.isfinite(current_portfolio_vol):
        return 0.0
    if target_vol <= 0:
        return 0.0
    if current_portfolio_vol <= target_vol:
        return 1.0
    return float(target_vol / current_portfolio_vol)


# ─────────────────────────────────────────────────────────────────────────────
# Pipeline completo Kelly + Vol Target
# ─────────────────────────────────────────────────────────────────────────────

def compute_kelly_vol_sizes(
    plan: list,
    equity: float,
    sigmas_daily: dict[str, float],
    payoff_b: float = DEFAULT_TP_SL_RATIO,
    kelly_fraction_mult: float = DEFAULT_KELLY_FRACTION,
    max_position_pct: float = DEFAULT_MAX_POSITION_PCT,
    target_vol_annual: float = DEFAULT_TARGET_VOL_ANNUAL,
    avg_correlation: float = DEFAULT_AVG_CORRELATION,
    correlations: Optional[np.ndarray] = None,
) -> dict[str, SizingResult]:
    """Orquesta Kelly fraccional + Vol Target sobre todo el plan.

    Args:
        plan: lista de objetos con atributos ticker, prob_win, side.
              Solo se sizean entries con side == 1 (LONG; Fase 2).
        equity: equity total de la cuenta en USD.
        sigmas_daily: dict {ticker -> sigma diaria estimada (e.g. EWM std).
                      Si falta un ticker, se usa DEFAULT_FALLBACK_SIGMA_DAILY.
        payoff_b: TP/SL ratio del Triple Barrier (default 2.0).
        kelly_fraction_mult: multiplicador (default 0.5 = Half-Kelly).
        max_position_pct: tope individual (default 0.25 = max 25% equity).
        target_vol_annual: vol anualizada objetivo del portfolio (default 0.20).
        avg_correlation: correlacion media usada si no tenemos matriz.
        correlations: matriz de correlaciones explicita (opcional).

    Returns:
        dict {ticker -> SizingResult}.
    """
    if equity <= 0:
        return {}

    # Paso 1: calcular kelly por ticker. Filtrar los con f* > 0 y side LONG.
    per_ticker_kelly: dict[str, tuple[float, float, float]] = {}
    for entry in plan:
        if getattr(entry, "side", 0) != 1:
            continue
        f_full, f_frac, f_cap = kelly_fraction(
            prob_win=float(entry.prob_win),
            payoff_b=payoff_b,
            kelly_fraction_mult=kelly_fraction_mult,
            cap=max_position_pct,
        )
        if f_cap > 0:
            per_ticker_kelly[entry.ticker] = (f_full, f_frac, f_cap)

    if not per_ticker_kelly:
        return {}

    # Paso 2: anualizar sigmas y construir vectores ordenados
    tickers = list(per_ticker_kelly.keys())
    sigmas_ann_list = []
    for t in tickers:
        sd = sigmas_daily.get(t, DEFAULT_FALLBACK_SIGMA_DAILY)
        sigmas_ann_list.append(annualized_vol(sd))

    kelly_caps_list = [per_ticker_kelly[t][2] for t in tickers]

    # Paso 3: estimar vol agregada del portfolio implicita con kelly_caps
    port_vol = portfolio_vol_estimate(
        weights=kelly_caps_list,
        sigmas_annual=sigmas_ann_list,
        correlations=correlations,
        avg_correlation=avg_correlation,
    )

    # Paso 4: scalar de vol targeting
    scalar = vol_target_scalar(port_vol, target_vol=target_vol_annual)

    # Paso 5: construir resultado por ticker
    out: dict[str, SizingResult] = {}
    for t, sigma_ann in zip(tickers, sigmas_ann_list):
        f_full, f_frac, f_cap = per_ticker_kelly[t]
        final_pct = f_cap * scalar
        note_parts = []
        if f_frac > f_cap:
            note_parts.append(f"capped_at_{max_position_pct*100:.0f}%")
        if scalar < 1.0:
            note_parts.append(f"vol_target_scaled_x{scalar:.2f}")
        # Para la prob_win, busca el entry original. Aqui solo guardamos
        # lo minimo que ya tenemos.
        prob = next(
            (float(e.prob_win) for e in plan
             if getattr(e, "ticker", None) == t and getattr(e, "side", 0) == 1),
            float("nan"),
        )
        out[t] = SizingResult(
            ticker=t,
            prob_win=prob,
            payoff_b=float(payoff_b),
            sigma_daily=float(sigmas_daily.get(t, DEFAULT_FALLBACK_SIGMA_DAILY)),
            sigma_annual=float(sigma_ann),
            kelly_full=round(f_full, 6),
            kelly_fractional=round(f_frac, 6),
            kelly_capped=round(f_cap, 6),
            portfolio_vol_estimate=round(float(port_vol), 6),
            vol_scalar=round(float(scalar), 6),
            final_pct=round(float(final_pct), 6),
            final_usd=round(float(final_pct * equity), 2),
            notes="|".join(note_parts) if note_parts else "ok",
        )
    return out


# ─────────────────────────────────────────────────────────────────────────────
# API de conveniencia: sizing por trade individual (sin plan agregado)
# ─────────────────────────────────────────────────────────────────────────────

def size_single_trade(
    prob_win: float,
    sigma_daily: float,
    equity: float,
    payoff_b: float = DEFAULT_TP_SL_RATIO,
    kelly_fraction_mult: float = DEFAULT_KELLY_FRACTION,
    max_position_pct: float = DEFAULT_MAX_POSITION_PCT,
    target_vol_annual: float = DEFAULT_TARGET_VOL_ANNUAL,
) -> SizingResult:
    """Sizing de un solo trade (cuando no hay otros operables).

    Sin portfolio implicito, vol target degenera: si la vol anual del activo
    excede target, recortamos el kelly proporcionalmente; si no, kelly puro.
    """
    f_full, f_frac, f_cap = kelly_fraction(prob_win, payoff_b,
                                           kelly_fraction_mult, max_position_pct)
    sigma_ann = annualized_vol(sigma_daily)
    # Vol del trade = peso * sigma_ann (aprox un solo activo)
    pos_vol = f_cap * sigma_ann
    scalar = vol_target_scalar(pos_vol, target_vol_annual)
    final_pct = f_cap * scalar
    notes = []
    if f_frac > f_cap:
        notes.append(f"capped_at_{max_position_pct*100:.0f}%")
    if scalar < 1.0:
        notes.append(f"vol_target_scaled_x{scalar:.2f}")
    return SizingResult(
        ticker="",
        prob_win=float(prob_win),
        payoff_b=float(payoff_b),
        sigma_daily=float(sigma_daily),
        sigma_annual=float(sigma_ann),
        kelly_full=round(f_full, 6),
        kelly_fractional=round(f_frac, 6),
        kelly_capped=round(f_cap, 6),
        portfolio_vol_estimate=round(float(pos_vol), 6),
        vol_scalar=round(float(scalar), 6),
        final_pct=round(float(final_pct), 6),
        final_usd=round(float(final_pct * max(equity, 0)), 2),
        notes="|".join(notes) if notes else "ok",
    )