"""Portfolio risk decomposition and optimization.

Cross-strategy portfolio risk: Euler decomposition, MCTR, risk budgeting,
mean-variance optimization, what-if analysis.

All functions are pure: data in, results out. No DB access, no API calls.
Requires numpy/scipy (core dependencies).

Usage::

    from cpz.risk.portfolio import mctr, euler_decomposition, optimize_portfolio

    contributions = euler_decomposition(weights, cov_matrix)
    marginal = mctr(weights, cov_matrix)
    optimal = optimize_portfolio(returns_matrix, method="risk_parity")
"""

from __future__ import annotations

import math
from typing import Any, Dict, List, Optional, Tuple

from pydantic import BaseModel, Field

import numpy as np
from scipy import optimize as sp_optimize

TRADING_DAYS_PER_YEAR: int = 252
EPSILON: float = 1e-15


# ── Models ────────────────────────────────────────────────────────────


class RiskContribution(BaseModel):
    """Per-strategy risk contribution from Euler decomposition."""
    strategy_id: str
    weight: float = 0.0
    marginal_contribution: float = 0.0
    risk_contribution: float = 0.0
    risk_contribution_pct: float = 0.0


class EulerDecomposition(BaseModel):
    """Euler risk decomposition result."""
    portfolio_volatility: float = 0.0
    contributions: List[RiskContribution] = Field(default_factory=list)


class MCTRResult(BaseModel):
    """Marginal contribution to risk per strategy."""
    portfolio_volatility: float = 0.0
    mctr: Dict[str, float] = Field(default_factory=dict)


class OptimizationResult(BaseModel):
    """Portfolio optimization result."""
    method: str = ""
    weights: Dict[str, float] = Field(default_factory=dict)
    expected_return: float = 0.0
    volatility: float = 0.0
    sharpe_ratio: float = 0.0
    max_drawdown_estimate: float = 0.0


class WhatIfResult(BaseModel):
    """What-if analysis for adding/changing a strategy."""
    current_volatility: float = 0.0
    new_volatility: float = 0.0
    volatility_change_pct: float = 0.0
    current_sharpe: float = 0.0
    new_sharpe: float = 0.0
    new_weights: Dict[str, float] = Field(default_factory=dict)
    diversification_benefit: float = 0.0


class RiskBudgetResult(BaseModel):
    """Risk budget check."""
    total_risk: float = 0.0
    budgets: Dict[str, float] = Field(default_factory=dict)
    actual: Dict[str, float] = Field(default_factory=dict)
    breaches: List[str] = Field(default_factory=list)
    within_budget: bool = True


# ── Public API ────────────────────────────────────────────────────────


def euler_decomposition(
    weights: Dict[str, float],
    cov_matrix: Dict[str, Dict[str, float]],
) -> EulerDecomposition:
    """Euler risk decomposition: contribution of each strategy to portfolio vol.

    The Euler decomposition satisfies: sum(contributions) = portfolio_volatility.
    Each contribution_i = w_i * (Sigma @ w)_i / sigma_p

    Args:
        weights: Strategy weights {strategy_id: weight}.
        cov_matrix: Covariance matrix {id_a: {id_b: cov_ab}}.

    Returns:
        EulerDecomposition with per-strategy risk contributions.
    """
    ids = list(weights.keys())
    n = len(ids)
    if n == 0:
        return EulerDecomposition()

    w = np.array([weights[i] for i in ids], dtype=np.float64)
    cov = _build_cov_matrix(ids, cov_matrix)

    sigma_p = float(np.sqrt(w @ cov @ w))
    if sigma_p < 1e-15:
        return EulerDecomposition(portfolio_volatility=0.0)

    # MCTR_i = (Sigma @ w)_i / sigma_p
    marginal = cov @ w / sigma_p

    # Risk contribution_i = w_i * MCTR_i
    rc = w * marginal
    total_rc = float(np.sum(rc))

    contributions = []
    for j, sid in enumerate(ids):
        contributions.append(RiskContribution(
            strategy_id=sid,
            weight=round(float(w[j]), 4),
            marginal_contribution=round(float(marginal[j]), 6),
            risk_contribution=round(float(rc[j]), 6),
            risk_contribution_pct=round(float(rc[j] / total_rc * 100), 2) if total_rc > 0 else 0.0,
        ))

    return EulerDecomposition(
        portfolio_volatility=round(sigma_p, 6),
        contributions=contributions,
    )


def mctr(
    weights: Dict[str, float],
    cov_matrix: Dict[str, Dict[str, float]],
) -> MCTRResult:
    """Marginal contribution to risk for each strategy.

    MCTR_i = (Sigma @ w)_i / sigma_p

    Args:
        weights: Strategy weights {strategy_id: weight}.
        cov_matrix: Covariance matrix {id_a: {id_b: cov_ab}}.

    Returns:
        MCTRResult with per-strategy marginal contribution.
    """
    ids = list(weights.keys())
    if not ids:
        return MCTRResult()

    w = np.array([weights[i] for i in ids], dtype=np.float64)
    cov = _build_cov_matrix(ids, cov_matrix)

    sigma_p = float(np.sqrt(w @ cov @ w))
    if sigma_p < 1e-15:
        return MCTRResult(portfolio_volatility=0.0)

    marginal = cov @ w / sigma_p

    return MCTRResult(
        portfolio_volatility=round(sigma_p, 6),
        mctr={sid: round(float(marginal[j]), 6) for j, sid in enumerate(ids)},
    )


def optimize_portfolio(
    returns_matrix: Dict[str, List[float]],
    method: str = "min_variance",
    risk_free_rate: float = 0.0,
    target_volatility: Optional[float] = None,
    *,
    trading_days: int = TRADING_DAYS_PER_YEAR,
    min_observations: int = 5,
    max_iter: int = 200,
    tolerance: float = 1e-12,
) -> OptimizationResult:
    """Portfolio optimization.

    Supported methods: min_variance, max_sharpe, risk_parity, equal_weight.

    Args:
        returns_matrix: Daily returns per strategy {strategy_id: [returns]}.
        method: Optimization method.
        risk_free_rate: Annual risk-free rate.
        target_volatility: Target annual volatility (optional constraint).
        trading_days: Trading days per year for annualisation (default 252).
        min_observations: Minimum data points required (default 5).
        max_iter: Maximum solver iterations (default 200).
        tolerance: Solver convergence tolerance (default 1e-12).

    Returns:
        OptimizationResult with optimal weights and expected metrics.
    """
    ids = list(returns_matrix.keys())
    n = len(ids)
    if n == 0:
        return OptimizationResult(method=method)

    # Build returns matrix and covariance
    min_len = min(len(returns_matrix[i]) for i in ids)
    if min_len < min_observations:
        return OptimizationResult(method=method)

    R = np.column_stack([np.array(returns_matrix[i][-min_len:], dtype=np.float64) for i in ids])
    mu = np.mean(R, axis=0) * trading_days
    cov = np.cov(R.T, ddof=1) * trading_days
    if cov.ndim == 0:
        cov = np.array([[float(cov)]])

    rf_daily = risk_free_rate / trading_days

    if method == "equal_weight":
        w = np.ones(n) / n
    elif method == "min_variance":
        w = _min_variance(cov, n)
    elif method == "max_sharpe":
        w = _max_sharpe(mu, cov, n, risk_free_rate)
    elif method == "risk_parity":
        w = _risk_parity(cov, n)
    else:
        w = np.ones(n) / n

    port_ret = float(w @ mu)
    port_vol = float(np.sqrt(w @ cov @ w))
    port_sharpe = (port_ret - risk_free_rate) / port_vol if port_vol > 0 else 0.0

    return OptimizationResult(
        method=method,
        weights={ids[j]: round(float(w[j]), 4) for j in range(n)},
        expected_return=round(port_ret * 100, 2),
        volatility=round(port_vol * 100, 2),
        sharpe_ratio=round(port_sharpe, 3),
    )


def what_if_analysis(
    current_weights: Dict[str, float],
    cov_matrix: Dict[str, Dict[str, float]],
    new_strategy_returns: List[float],
    new_strategy_id: str,
    allocation_pct: float = 10.0,
    existing_returns: Optional[Dict[str, List[float]]] = None,
    risk_free_rate: float = 0.0,
) -> WhatIfResult:
    """What-if analysis: what happens to portfolio risk if a new strategy is added?

    Args:
        current_weights: Current strategy weights.
        cov_matrix: Current covariance matrix.
        new_strategy_returns: Daily returns of the new strategy.
        new_strategy_id: ID for the new strategy.
        allocation_pct: Percentage to allocate to the new strategy.
        existing_returns: Current strategy returns (for Sharpe calculation).
        risk_free_rate: Annual risk-free rate.

    Returns:
        WhatIfResult with before/after risk comparison.
    """
    ids = list(current_weights.keys())
    if not ids:
        return WhatIfResult()

    w_old = np.array([current_weights[i] for i in ids], dtype=np.float64)
    cov_old = _build_cov_matrix(ids, cov_matrix)
    sigma_old = float(np.sqrt(w_old @ cov_old @ w_old))

    # Scale down existing weights and add new strategy
    alloc = allocation_pct / 100.0
    scale = 1.0 - alloc
    new_ids = ids + [new_strategy_id]
    w_new = np.append(w_old * scale, alloc)

    # Extend covariance matrix with the new strategy
    new_r = np.array(new_strategy_returns, dtype=np.float64)
    n_new = len(new_ids)
    cov_new = np.zeros((n_new, n_new))
    cov_new[:n_new - 1, :n_new - 1] = cov_old

    # Estimate new covariances (if existing returns provided)
    if existing_returns:
        for j, sid in enumerate(ids):
            if sid in existing_returns:
                er = np.array(existing_returns[sid], dtype=np.float64)
                mn = min(len(er), len(new_r))
                if mn > 5:
                    c = float(np.cov(er[-mn:], new_r[-mn:], ddof=1)[0, 1])
                    cov_new[j, n_new - 1] = c
                    cov_new[n_new - 1, j] = c
    cov_new[n_new - 1, n_new - 1] = float(np.var(new_r, ddof=1)) if len(new_r) > 1 else 0.0

    sigma_new = float(np.sqrt(w_new @ cov_new @ w_new))
    vol_change = ((sigma_new - sigma_old) / sigma_old * 100) if sigma_old > 0 else 0.0

    # Sharpe comparison
    sharpe_old = 0.0
    sharpe_new = 0.0
    if existing_returns:
        port_ret_old = sum(current_weights.get(sid, 0) * np.mean(existing_returns.get(sid, [0])) * 252
                          for sid in ids)
        sharpe_old = (port_ret_old - risk_free_rate) / (sigma_old * math.sqrt(252)) if sigma_old > 0 else 0.0
        port_ret_new = port_ret_old * scale + alloc * float(np.mean(new_r)) * 252
        sharpe_new = (port_ret_new - risk_free_rate) / (sigma_new * math.sqrt(252)) if sigma_new > 0 else 0.0

    diversification = max(0, sigma_old - sigma_new)

    return WhatIfResult(
        current_volatility=round(sigma_old * math.sqrt(252) * 100, 2),
        new_volatility=round(sigma_new * math.sqrt(252) * 100, 2),
        volatility_change_pct=round(vol_change, 2),
        current_sharpe=round(sharpe_old, 3),
        new_sharpe=round(sharpe_new, 3),
        new_weights={new_ids[j]: round(float(w_new[j]), 4) for j in range(n_new)},
        diversification_benefit=round(diversification * math.sqrt(252) * 100, 2),
    )


def risk_budget_check(
    weights: Dict[str, float],
    cov_matrix: Dict[str, Dict[str, float]],
    budgets: Dict[str, float],
) -> RiskBudgetResult:
    """Check if strategies are within their risk budgets.

    Args:
        weights: Strategy weights.
        cov_matrix: Covariance matrix.
        budgets: Per-strategy risk budget as percentage of total risk.

    Returns:
        RiskBudgetResult with actual vs budgeted risk and breaches.
    """
    decomp = euler_decomposition(weights, cov_matrix)
    actual = {c.strategy_id: c.risk_contribution_pct for c in decomp.contributions}

    breaches = []
    for sid, budget_pct in budgets.items():
        actual_pct = actual.get(sid, 0.0)
        if actual_pct > budget_pct:
            breaches.append(sid)

    return RiskBudgetResult(
        total_risk=decomp.portfolio_volatility,
        budgets=budgets,
        actual=actual,
        breaches=breaches,
        within_budget=len(breaches) == 0,
    )


# ── Optimization internals ───────────────────────────────────────────


def _build_cov_matrix(ids: List[str], cov_dict: Dict[str, Dict[str, float]]) -> "np.ndarray":
    """Convert dict-of-dicts covariance to numpy matrix."""
    n = len(ids)
    cov = np.zeros((n, n))
    for i, a in enumerate(ids):
        for j, b in enumerate(ids):
            cov[i, j] = cov_dict.get(a, {}).get(b, 0.0)
    return cov


def _min_variance(cov: "np.ndarray", n: int) -> "np.ndarray":
    """Minimum variance portfolio via quadratic programming."""
    try:
        inv_cov = np.linalg.inv(cov)
        ones = np.ones(n)
        w = inv_cov @ ones / (ones @ inv_cov @ ones)
        return np.clip(w, 0, 1)
    except np.linalg.LinAlgError:
        return np.ones(n) / n


def _max_sharpe(mu: "np.ndarray", cov: "np.ndarray", n: int, rf: float) -> "np.ndarray":
    """Maximum Sharpe ratio portfolio."""
    try:
        excess = mu - rf
        inv_cov = np.linalg.inv(cov)
        w = inv_cov @ excess
        w_sum = np.sum(w)
        if w_sum > 0:
            w = w / w_sum
        else:
            w = np.ones(n) / n
        return np.clip(w, 0, 1)
    except np.linalg.LinAlgError:
        return np.ones(n) / n


def _risk_parity(cov: "np.ndarray", n: int) -> "np.ndarray":
    """Risk parity: equal risk contribution from each strategy.

    Uses iterative optimization to find weights where each strategy
    contributes equally to total portfolio risk.
    """
    def objective(w):
        w = np.abs(w)
        sigma = np.sqrt(w @ cov @ w)
        if sigma < 1e-15:
            return 1e10
        rc = w * (cov @ w) / sigma
        target = sigma / n
        return float(np.sum((rc - target) ** 2))

    x0 = np.ones(n) / n
    bounds = [(0.001, 1.0)] * n
    constraints = {"type": "eq", "fun": lambda w: np.sum(w) - 1.0}

    try:
        result = sp_optimize.minimize(
            objective, x0, method="SLSQP",
            bounds=bounds, constraints=constraints,
            options={"maxiter": max_iter, "ftol": tolerance},
        )
        w = np.abs(result.x)
        return w / np.sum(w)
    except Exception:
        return np.ones(n) / n
