"""Risk management module for jbqlab.

This module provides risk management utilities including:
- Stop loss calculations
- Position limits
- Risk budgeting
- Drawdown controls
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    pass

__all__ = [
    "RiskLimits",
    "calculate_stop_loss",
    "calculate_take_profit",
    "apply_risk_limits",
    "calculate_risk_budget",
    "drawdown_control",
    "volatility_scaling",
]


@dataclass
class RiskLimits:
    """Risk management limits configuration."""

    max_position_size: float = 1.0  # Max fraction of capital
    max_drawdown: float = 0.20  # 20% max drawdown
    max_daily_loss: float = 0.05  # 5% max daily loss
    max_volatility: float = 0.30  # 30% annualized vol
    stop_loss_pct: float | None = None
    take_profit_pct: float | None = None
    trailing_stop_pct: float | None = None

    def __post_init__(self):
        """Validate limits."""
        if not 0 < self.max_position_size <= 1:
            raise ValueError("max_position_size must be between 0 and 1")
        if not 0 < self.max_drawdown <= 1:
            raise ValueError("max_drawdown must be between 0 and 1")


def calculate_stop_loss(
    entry_price: float,
    method: str = "fixed",
    pct: float = 0.02,
    atr_value: float | None = None,
    atr_mult: float = 2.0,
) -> float:
    """Calculate stop loss price.

    Args:
        entry_price: Entry price of position.
        method: Stop method ('fixed', 'atr').
        pct: Fixed percentage for stop (default 2%).
        atr_value: ATR value for ATR-based stop.
        atr_mult: ATR multiplier.

    Returns:
        Stop loss price.
    """
    if method == "fixed":
        return entry_price * (1 - pct)
    elif method == "atr":
        if atr_value is None:
            raise ValueError("ATR value required for ATR-based stop")
        return entry_price - (atr_mult * atr_value)
    else:
        raise ValueError(f"Invalid method: {method}")


def calculate_take_profit(
    entry_price: float,
    risk_reward: float = 2.0,
    stop_loss_price: float | None = None,
    fixed_pct: float | None = None,
) -> float:
    """Calculate take profit price.

    Args:
        entry_price: Entry price.
        risk_reward: Risk/reward ratio (default 2:1).
        stop_loss_price: Stop loss price (for R:R calculation).
        fixed_pct: Fixed percentage target.

    Returns:
        Take profit price.
    """
    if fixed_pct is not None:
        return entry_price * (1 + fixed_pct)

    if stop_loss_price is not None:
        risk = entry_price - stop_loss_price
        return entry_price + (risk * risk_reward)

    raise ValueError("Must provide stop_loss_price or fixed_pct")


def apply_risk_limits(
    signal: pd.Series,
    prices: pd.Series,
    limits: RiskLimits,
) -> pd.Series:
    """Apply risk limits to a signal series.

    Modifies signals based on risk limits (drawdown, daily loss, etc.).

    Args:
        signal: Original signal series.
        prices: Price series.
        limits: RiskLimits configuration.

    Returns:
        Modified signal with risk limits applied.
    """
    modified = signal.copy()
    equity = pd.Series(index=prices.index, dtype=float)

    position = 0
    equity.iloc[0] = 100000  # Starting capital
    peak_equity = equity.iloc[0]
    entry_price = None

    for i in range(1, len(prices)):
        # Calculate equity
        if position == 1:
            daily_return = (prices.iloc[i] / prices.iloc[i-1]) - 1
            equity.iloc[i] = equity.iloc[i-1] * (1 + daily_return)
        else:
            equity.iloc[i] = equity.iloc[i-1]

        # Update peak
        if equity.iloc[i] > peak_equity:
            peak_equity = equity.iloc[i]

        # Current drawdown
        current_dd = (peak_equity - equity.iloc[i]) / peak_equity

        # Daily loss
        daily_loss = 1 - (equity.iloc[i] / equity.iloc[i-1]) if i > 0 else 0

        # Check limits
        should_exit = False

        # Max drawdown breach
        if current_dd >= limits.max_drawdown:
            should_exit = True

        # Max daily loss breach
        if daily_loss >= limits.max_daily_loss:
            should_exit = True

        # Stop loss check
        if position == 1 and limits.stop_loss_pct and entry_price:
            stop_price = entry_price * (1 - limits.stop_loss_pct)
            if prices.iloc[i] <= stop_price:
                should_exit = True

        # Take profit check
        if position == 1 and limits.take_profit_pct and entry_price:
            tp_price = entry_price * (1 + limits.take_profit_pct)
            if prices.iloc[i] >= tp_price:
                should_exit = True

        # Apply exit if needed
        if should_exit:
            modified.iloc[i] = 0

        # Track position state
        if modified.iloc[i] == 1 and position == 0:
            position = 1
            entry_price = prices.iloc[i]
        elif modified.iloc[i] == 0 and position == 1:
            position = 0
            entry_price = None

    return modified


def calculate_risk_budget(
    total_capital: float,
    max_risk_per_trade: float,
    n_positions: int,
    correlation: float = 0.5,
) -> dict:
    """Calculate risk budget for portfolio.

    Args:
        total_capital: Total portfolio capital.
        max_risk_per_trade: Maximum risk per trade (fraction).
        n_positions: Number of positions.
        correlation: Average correlation between positions.

    Returns:
        Dictionary with risk budget metrics.
    """
    # Individual position risk
    individual_risk = total_capital * max_risk_per_trade

    # Diversification factor
    div_factor = np.sqrt(n_positions + correlation * n_positions * (n_positions - 1))

    # Portfolio risk
    portfolio_risk = individual_risk * div_factor

    # Risk per position (accounting for correlation)
    adjusted_risk_per_position = portfolio_risk / n_positions

    return {
        "total_capital": total_capital,
        "individual_risk": individual_risk,
        "n_positions": n_positions,
        "correlation_assumption": correlation,
        "diversification_factor": div_factor,
        "portfolio_risk": portfolio_risk,
        "adjusted_risk_per_position": adjusted_risk_per_position,
        "capital_per_position": total_capital / n_positions,
    }


def drawdown_control(
    signal: pd.Series,
    equity: pd.Series,
    max_drawdown: float = 0.15,
    recovery_threshold: float = 0.05,
) -> pd.Series:
    """Apply drawdown-based position sizing.

    Reduces or eliminates positions when drawdown exceeds threshold.
    Gradually recovers as drawdown improves.

    Args:
        signal: Original signal.
        equity: Equity curve.
        max_drawdown: Maximum allowed drawdown.
        recovery_threshold: Drawdown level to start recovery.

    Returns:
        Modified signal with drawdown control.
    """
    modified = signal.copy()

    peak = equity.cummax()
    drawdown = (peak - equity) / peak

    # Position multiplier based on drawdown
    multiplier = pd.Series(1.0, index=equity.index)

    for i in range(len(equity)):
        dd = drawdown.iloc[i]

        if dd >= max_drawdown:
            # Full exit
            multiplier.iloc[i] = 0
        elif dd >= recovery_threshold:
            # Linear scale down
            scale = 1 - (dd - recovery_threshold) / (max_drawdown - recovery_threshold)
            multiplier.iloc[i] = max(0, scale)

    # Apply multiplier (convert to binary for simplicity)
    modified = (signal * (multiplier >= 0.5)).astype(int)

    return modified


def volatility_scaling(
    signal: pd.Series,
    returns: pd.Series,
    target_vol: float = 0.15,
    lookback: int = 20,
    min_scale: float = 0.5,
    max_scale: float = 2.0,
) -> pd.Series:
    """Scale position size based on volatility.

    When volatility is high, reduce position. When low, increase.

    Args:
        signal: Original signal.
        returns: Return series for volatility calculation.
        target_vol: Target annualized volatility.
        lookback: Lookback window for volatility.
        min_scale: Minimum scale factor.
        max_scale: Maximum scale factor.

    Returns:
        Volatility-scaled signal (can be fractional).
    """
    # Rolling volatility (annualized)
    rolling_vol = returns.rolling(window=lookback).std() * np.sqrt(252)

    # Scale factor
    scale = target_vol / rolling_vol.clip(lower=0.01)  # Avoid division by zero
    scale = scale.clip(lower=min_scale, upper=max_scale)
    scale = scale.fillna(1.0)

    # Apply scaling
    scaled_signal = signal * scale

    return scaled_signal


def max_consecutive_losses(returns: pd.Series) -> int:
    """Calculate maximum consecutive losing days.

    Args:
        returns: Return series.

    Returns:
        Maximum number of consecutive losses.
    """
    is_loss = returns < 0
    groups = (is_loss != is_loss.shift()).cumsum()
    consecutive = is_loss.groupby(groups).cumsum()
    return int(consecutive.max())


def ulcer_index(equity: pd.Series, window: int = 14) -> pd.Series:
    """Calculate Ulcer Index (pain of drawdowns).

    Lower is better. Considers both depth and duration of drawdowns.

    Args:
        equity: Equity curve.
        window: Rolling window.

    Returns:
        Ulcer Index series.
    """
    peak = equity.rolling(window=window, min_periods=1).max()
    pct_drawdown = ((equity - peak) / peak) * 100

    squared_dd = pct_drawdown ** 2
    ulcer = np.sqrt(squared_dd.rolling(window=window).mean())

    return ulcer
