"""Markov-driven regime-blended risk measure for portfolio optimization.

Implements the probability-weighted risk objective::

    ρ_t(w) = Σ_s  γ_T(s) · ρ_s(w)

where γ_T(s) = P(z_T = s | r_{1:T}) is the current filtered state probability
from a fitted HMM, and ρ_s is the regime-specific risk measure for state s.

Also implements regime-conditional risk budgets::

    b_t = Σ_s γ_T(s) · b_s

where b_s is the budget vector for regime s, blended by current state probs.

Usage example::

    from optimizer.moments._hmm import fit_hmm, HMMConfig
    from optimizer.optimization._regime_risk import (
        RegimeRiskConfig,
        compute_blended_risk_measure,
        build_regime_blended_optimizer,
        compute_regime_budget,
        build_regime_risk_budgeting,
    )

    hmm_result = fit_hmm(returns, HMMConfig(n_states=2))

    config = RegimeRiskConfig.for_calm_stress()
    risk = compute_blended_risk_measure(
        returns, weights, hmm_result, config.regime_measures
    )
    optimizer = build_regime_blended_optimizer(config, hmm_result)
"""

from __future__ import annotations

import logging
import warnings
from collections.abc import Sequence
from dataclasses import dataclass, field

import numpy as np
import numpy.typing as npt
import pandas as pd
from skfolio.optimization import MeanRisk, RiskBudgeting

from optimizer.exceptions import ConfigurationError, DataError
from optimizer.moments._hmm import HMMConfig, HMMResult
from optimizer.optimization._config import RiskBudgetingConfig, RiskMeasureType
from optimizer.optimization._factory import _RISK_MEASURE_MAP, build_risk_budgeting

logger = logging.getLogger(__name__)

# Minimum per-regime observations before falling back to full-sample risk.
_MIN_REGIME_SAMPLES: int = 5


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class RegimeRiskConfig:
    """Configuration for Markov-driven blended risk measure optimisation.

    Attributes
    ----------
    regime_measures : tuple[RiskMeasureType, ...]
        One :class:`RiskMeasureType` per HMM state (index = state label).
        E.g. ``(VARIANCE, CVAR)`` means state-0 uses variance and state-1
        uses CVaR.  Must match ``HMMResult.n_states``.
    hmm_config : HMMConfig
        HMM hyper-parameters used when fitting inside the pipeline.
    cvar_beta : float
        Confidence level for CVaR / ES computation (default 0.95 = 95th
        percentile).
    """

    regime_measures: tuple[RiskMeasureType, ...]
    hmm_config: HMMConfig = field(default_factory=HMMConfig)
    cvar_beta: float = 0.95

    @classmethod
    def for_calm_stress(cls, **kwargs: object) -> RegimeRiskConfig:
        """Two-regime preset: calm → variance, stress → CVaR."""
        return cls(
            regime_measures=(RiskMeasureType.VARIANCE, RiskMeasureType.CVAR),
            hmm_config=HMMConfig(n_states=2),
            **kwargs,  # type: ignore[arg-type]
        )

    @classmethod
    def for_calm_stress_drawdown(cls, **kwargs: object) -> RegimeRiskConfig:
        """Two-regime preset: calm → variance, stress → CDaR."""
        return cls(
            regime_measures=(RiskMeasureType.VARIANCE, RiskMeasureType.CDAR),
            hmm_config=HMMConfig(n_states=2),
            cvar_beta=0.95,
            **kwargs,
        )

    @classmethod
    def for_three_regimes(cls, **kwargs: object) -> RegimeRiskConfig:
        """Three-regime preset: calm → variance, normal → MAD, stress → CVaR."""
        return cls(
            regime_measures=(
                RiskMeasureType.VARIANCE,
                RiskMeasureType.MEAN_ABSOLUTE_DEVIATION,
                RiskMeasureType.CVAR,
            ),
            hmm_config=HMMConfig(n_states=3),
            **kwargs,  # type: ignore[arg-type]
        )


# ---------------------------------------------------------------------------
# Per-regime risk computation (pure, numpy-only, no skfolio fitting needed)
# ---------------------------------------------------------------------------


def _compute_regime_risk(
    portfolio_returns: npt.NDArray[np.float64],
    measure: RiskMeasureType,
    cvar_beta: float = 0.95,
) -> float:
    """Compute a scalar risk for *portfolio_returns* using *measure*.

    Supports: VARIANCE, STANDARD_DEVIATION, SEMI_VARIANCE, SEMI_DEVIATION,
    MEAN_ABSOLUTE_DEVIATION, CVAR, WORST_REALIZATION.  All other measures
    fall back to standard deviation.

    Returns 0.0 for an empty array.
    """
    n = len(portfolio_returns)
    if n == 0:
        return 0.0

    r = portfolio_returns
    ddof = 1 if n > 1 else 0

    if measure == RiskMeasureType.VARIANCE:
        return float(np.var(r, ddof=ddof))

    if measure == RiskMeasureType.STANDARD_DEVIATION:
        return float(np.std(r, ddof=ddof))

    if measure == RiskMeasureType.SEMI_VARIANCE:
        mu = np.mean(r)
        downside = r[r < mu] - mu
        if len(downside) < 2:
            return 0.0
        return float(np.mean(downside**2))

    if measure == RiskMeasureType.SEMI_DEVIATION:
        mu = np.mean(r)
        downside = r[r < mu] - mu
        if len(downside) < 2:
            return 0.0
        return float(np.sqrt(np.mean(downside**2)))

    if measure == RiskMeasureType.MEAN_ABSOLUTE_DEVIATION:
        return float(np.mean(np.abs(r - np.mean(r))))

    if measure == RiskMeasureType.CVAR:
        sorted_r = np.sort(r)
        cutoff = max(1, int(np.floor((1.0 - cvar_beta) * n)))
        return float(-np.mean(sorted_r[:cutoff]))

    if measure == RiskMeasureType.WORST_REALIZATION:
        return max(float(-np.min(r)), 0.0)

    # -- Drawdown-based measures --
    # Compute cumulative wealth and drawdown series
    if measure in (
        RiskMeasureType.MAX_DRAWDOWN,
        RiskMeasureType.AVERAGE_DRAWDOWN,
        RiskMeasureType.CDAR,
        RiskMeasureType.ULCER_INDEX,
    ):
        cumulative = np.cumprod(1.0 + r)
        running_max = np.maximum.accumulate(cumulative)
        drawdown = 1.0 - cumulative / running_max

        if measure == RiskMeasureType.MAX_DRAWDOWN:
            return float(np.max(drawdown))

        if measure == RiskMeasureType.AVERAGE_DRAWDOWN:
            return float(np.mean(drawdown))

        if measure == RiskMeasureType.ULCER_INDEX:
            return float(np.sqrt(np.mean(drawdown**2)))

        if measure == RiskMeasureType.CDAR:
            sorted_dd = np.sort(drawdown)[::-1]
            cutoff = max(1, int(np.floor((1.0 - cvar_beta) * n)))
            return float(np.mean(sorted_dd[:cutoff]))

    if measure == RiskMeasureType.EDAR:
        raise NotImplementedError(
            "EDAR (Entropic Drawdown at Risk) is not implemented for "
            "regime-conditional risk computation"
        )

    warnings.warn(
        f"Risk measure {measure.value!r} is not supported for "
        "regime-conditional risk computation; falling back to standard deviation",
        UserWarning,
        stacklevel=2,
    )
    return float(np.std(r, ddof=ddof))


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def compute_blended_risk_measure(
    returns: pd.DataFrame,
    weights: npt.NDArray[np.float64],
    hmm_result: HMMResult,
    regime_measures: Sequence[RiskMeasureType],
    cvar_beta: float = 0.95,
) -> float:
    """Compute ρ_t(w) = Σ_s γ_T(s) · ρ_s(w): probability-weighted blended risk.

    Each regime's risk ρ_s(w) is computed on the historical portfolio returns
    from periods assigned to that regime (hard Viterbi-style assignment via
    argmax of filtered probabilities).  Regimes with fewer than
    ``_MIN_REGIME_SAMPLES`` observations fall back to full-sample risk.

    The blend weights γ_T(s) are the **last row** of
    ``hmm_result.filtered_probs`` — the most-recent filtered state
    probabilities, which sum to 1 by HMM construction.

    Args:
        returns: Dates × assets matrix of linear returns.
        weights: 1-D array of portfolio weights, shape ``(n_assets,)``.
        hmm_result: Fitted HMM with ``filtered_probs`` aligned to *returns*.
        regime_measures: One :class:`RiskMeasureType` per HMM state.
        cvar_beta: CVaR confidence level (used when a regime uses CVaR).

    Returns:
        Scalar non-negative blended risk value.

    Raises:
        ValueError: If ``len(regime_measures) != hmm_result.n_states``.
        ValueError: If *weights* length does not match *returns* columns.
    """
    n_states = hmm_result.filtered_probs.shape[1]
    if len(regime_measures) != n_states:
        raise ConfigurationError(
            f"len(regime_measures) = {len(regime_measures)} must match "
            f"hmm_result n_states = {n_states}."
        )

    n_assets = returns.shape[1]
    if len(weights) != n_assets:
        raise DataError(f"weights length {len(weights)} != returns columns {n_assets}.")

    # Current regime probabilities (last timestep of the filtered sequence)
    gamma: npt.NDArray[np.float64] = hmm_result.filtered_probs.iloc[-1].to_numpy(
        dtype=np.float64
    )

    # Align on common dates
    common_idx = returns.index.intersection(hmm_result.filtered_probs.index)
    if len(common_idx) == 0:
        return 0.0

    aligned_returns = returns.loc[common_idx].to_numpy(dtype=np.float64)
    aligned_probs = hmm_result.filtered_probs.loc[common_idx].to_numpy(dtype=np.float64)

    portfolio_returns: npt.NDArray[np.float64] = aligned_returns @ weights
    regime_assignments: npt.NDArray[np.intp] = np.argmax(aligned_probs, axis=1)

    blended = 0.0
    for s in range(n_states):
        if gamma[s] == 0.0:
            continue

        mask = regime_assignments == s
        n_regime = int(mask.sum())

        p_ret = (
            portfolio_returns[mask]
            if n_regime >= _MIN_REGIME_SAMPLES
            else portfolio_returns  # fall back to full sample
        )

        risk_s = _compute_regime_risk(p_ret, regime_measures[s], cvar_beta)
        blended += gamma[s] * risk_s

    return float(blended)


def build_regime_blended_optimizer(
    config: RegimeRiskConfig,
    hmm_result: HMMResult,
    **mean_risk_kwargs: object,
) -> MeanRisk:
    """Return a MeanRisk optimizer configured for the dominant regime.

    Because skfolio's ``MeanRisk`` requires a single convex risk measure,
    this function selects the risk measure of the **dominant regime** —
    the state s* = argmax_s γ_T(s) — and builds a standard
    :class:`skfolio.optimization.MeanRisk` with that measure.

    Args:
        config: Regime risk configuration.
        hmm_result: Fitted HMM providing current regime probabilities.
        **mean_risk_kwargs: Additional keyword arguments forwarded to
            :class:`skfolio.optimization.MeanRisk` (e.g.
            ``min_weights``, ``max_weights``, ``prior_estimator``).

    Returns:
        :class:`skfolio.optimization.MeanRisk` instance ready for
        ``fit(X)`` calls.

    Raises:
        ValueError: If ``len(config.regime_measures) != hmm_result.n_states``.
    """
    n_states = hmm_result.filtered_probs.shape[1]
    if len(config.regime_measures) != n_states:
        raise ConfigurationError(
            f"len(config.regime_measures) = {len(config.regime_measures)} must "
            f"match hmm_result n_states = {n_states}."
        )

    gamma: npt.NDArray[np.float64] = hmm_result.filtered_probs.iloc[-1].to_numpy(
        dtype=np.float64
    )
    dominant_state = int(np.argmax(gamma))
    dominant_measure = config.regime_measures[dominant_state]
    skfolio_measure = _RISK_MEASURE_MAP[dominant_measure]

    return MeanRisk(risk_measure=skfolio_measure, **mean_risk_kwargs)


# ---------------------------------------------------------------------------
# Regime-conditional risk budgets
# ---------------------------------------------------------------------------


def compute_regime_budget(
    regime_budgets: list[npt.NDArray[np.float64]],
    filtered_probs_last: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]:
    """Compute b_t = Σ_s γ_T(s) · b_s: probability-weighted blended budget.

    Args:
        regime_budgets: One budget vector per HMM state, each shape
            ``(n_assets,)``.  Values must be non-negative.
        filtered_probs_last: Current state probabilities γ_T, shape
            ``(n_states,)``.  Typically the last row of
            ``HMMResult.filtered_probs``.

    Returns:
        Blended budget vector of shape ``(n_assets,)``, normalised to sum
        to 1.

    Raises:
        ValueError: If ``len(regime_budgets) != len(filtered_probs_last)``.
    """
    n_states = len(filtered_probs_last)
    if len(regime_budgets) != n_states:
        raise ConfigurationError(
            f"len(regime_budgets) = {len(regime_budgets)} must match "
            f"len(filtered_probs_last) = {n_states}."
        )

    gamma = filtered_probs_last.astype(np.float64)
    blended: npt.NDArray[np.float64] = np.zeros(
        len(regime_budgets[0]), dtype=np.float64
    )
    for s in range(n_states):
        blended += gamma[s] * regime_budgets[s].astype(np.float64)

    total = float(blended.sum())
    if total <= 0.0:
        n_assets = len(blended)
        return np.full(n_assets, 1.0 / n_assets, dtype=np.float64)
    return blended / total


def build_regime_risk_budgeting(
    config: RiskBudgetingConfig,
    hmm_result: HMMResult,
    regime_budgets: list[npt.NDArray[np.float64]],
    **kwargs: object,
) -> RiskBudgeting:
    """Return a RiskBudgeting optimizer with a regime-blended budget vector.

    Computes ``b_t = Σ_s γ_T(s) · b_s`` from the current HMM state
    probabilities and injects the result as ``risk_budget`` into
    :func:`build_risk_budgeting`.

    Args:
        config: Risk-budgeting configuration (risk measure, constraints, …).
        hmm_result: Fitted HMM providing current filtered state probabilities.
        regime_budgets: One budget vector per HMM state, each shape
            ``(n_assets,)``.
        **kwargs: Additional keyword arguments forwarded to
            :func:`build_risk_budgeting` (e.g. ``prior_estimator``).

    Returns:
        :class:`skfolio.optimization.RiskBudgeting` instance ready for
        ``fit(X)`` calls.

    Raises:
        ValueError: If ``len(regime_budgets) != hmm_result.n_states``.
    """
    gamma: npt.NDArray[np.float64] = hmm_result.filtered_probs.iloc[-1].to_numpy(
        dtype=np.float64
    )
    blended_budget = compute_regime_budget(regime_budgets, gamma)
    return build_risk_budgeting(config, risk_budget=blended_budget, **kwargs)
