"""Portfolio analysis module for jbqlab.

This module provides tools for analyzing portfolio performance,
including position sizing, risk metrics, and allocation analysis.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    from collections.abc import Mapping

__all__ = [
    "PortfolioStats",
    "calculate_portfolio_stats",
    "position_sizing",
    "kelly_criterion",
    "risk_parity_weights",
    "correlation_matrix",
    "rolling_correlation",
    "diversification_ratio",
]


@dataclass
class PortfolioStats:
    """Portfolio statistics container."""

    total_return: float
    cagr: float
    volatility: float
    sharpe_ratio: float
    sortino_ratio: float
    max_drawdown: float
    var_95: float  # Value at Risk (95%)
    cvar_95: float  # Conditional VaR (Expected Shortfall)
    win_rate: float
    profit_factor: float
    avg_win: float
    avg_loss: float
    best_day: float
    worst_day: float
    skewness: float
    kurtosis: float

    def to_dict(self) -> dict[str, float]:
        """Convert to dictionary."""
        return {
            "total_return": self.total_return,
            "cagr": self.cagr,
            "volatility": self.volatility,
            "sharpe_ratio": self.sharpe_ratio,
            "sortino_ratio": self.sortino_ratio,
            "max_drawdown": self.max_drawdown,
            "var_95": self.var_95,
            "cvar_95": self.cvar_95,
            "win_rate": self.win_rate,
            "profit_factor": self.profit_factor,
            "avg_win": self.avg_win,
            "avg_loss": self.avg_loss,
            "best_day": self.best_day,
            "worst_day": self.worst_day,
            "skewness": self.skewness,
            "kurtosis": self.kurtosis,
        }


def calculate_portfolio_stats(
    equity_curve: pd.Series,
    risk_free_rate: float = 0.0,
    periods_per_year: int = 252,
) -> PortfolioStats:
    """Calculate comprehensive portfolio statistics.

    Args:
        equity_curve: Series of portfolio values over time.
        risk_free_rate: Annual risk-free rate (default 0).
        periods_per_year: Trading periods per year (default 252).

    Returns:
        PortfolioStats with comprehensive metrics.
    """
    returns = equity_curve.pct_change().dropna()

    # Basic metrics
    total_return = (equity_curve.iloc[-1] / equity_curve.iloc[0]) - 1
    n_years = len(equity_curve) / periods_per_year
    cagr = (1 + total_return) ** (1 / n_years) - 1 if n_years > 0 else 0

    volatility = returns.std() * np.sqrt(periods_per_year)

    # Sharpe and Sortino
    excess_returns = returns - risk_free_rate / periods_per_year
    sharpe = (
        excess_returns.mean() / returns.std() * np.sqrt(periods_per_year)
        if returns.std() > 0
        else 0
    )

    downside_returns = returns[returns < 0]
    downside_std = downside_returns.std() if len(downside_returns) > 0 else 0
    sortino = (
        excess_returns.mean() / downside_std * np.sqrt(periods_per_year)
        if downside_std > 0
        else 0
    )

    # Drawdown
    rolling_max = equity_curve.cummax()
    drawdown = (equity_curve - rolling_max) / rolling_max
    max_drawdown = drawdown.min()

    # VaR and CVaR (95%)
    var_95 = returns.quantile(0.05)
    cvar_95 = returns[returns <= var_95].mean() if (returns <= var_95).any() else var_95

    # Win/Loss metrics
    positive_returns = returns[returns > 0]
    negative_returns = returns[returns < 0]

    win_rate = len(positive_returns) / len(returns) if len(returns) > 0 else 0
    avg_win = positive_returns.mean() if len(positive_returns) > 0 else 0
    avg_loss = negative_returns.mean() if len(negative_returns) > 0 else 0

    gross_profit = positive_returns.sum()
    gross_loss = abs(negative_returns.sum())
    profit_factor = gross_profit / gross_loss if gross_loss > 0 else float("inf")

    # Best/worst
    best_day = returns.max()
    worst_day = returns.min()

    # Higher moments
    skewness = returns.skew()
    kurtosis = returns.kurtosis()

    return PortfolioStats(
        total_return=total_return,
        cagr=cagr,
        volatility=volatility,
        sharpe_ratio=sharpe,
        sortino_ratio=sortino,
        max_drawdown=max_drawdown,
        var_95=var_95,
        cvar_95=cvar_95,
        win_rate=win_rate,
        profit_factor=profit_factor,
        avg_win=avg_win,
        avg_loss=avg_loss,
        best_day=best_day,
        worst_day=worst_day,
        skewness=skewness,
        kurtosis=kurtosis,
    )


def position_sizing(
    method: str,
    capital: float,
    price: float,
    volatility: float | None = None,
    risk_per_trade: float = 0.02,
    stop_loss_pct: float | None = None,
    **kwargs: Any,
) -> int:
    """Calculate position size based on various methods.

    Args:
        method: Sizing method ('fixed', 'volatility', 'risk', 'equal').
        capital: Available capital.
        price: Asset price.
        volatility: Asset volatility (required for 'volatility' method).
        risk_per_trade: Maximum risk per trade as fraction of capital.
        stop_loss_pct: Stop loss percentage (required for 'risk' method).
        **kwargs: Additional method-specific parameters.

    Returns:
        Number of shares/units to buy.

    Raises:
        ValueError: If invalid method or missing required parameters.
    """
    if method == "fixed":
        # Fixed percentage of capital
        pct = kwargs.get("allocation_pct", 0.1)
        position_value = capital * pct
        return int(position_value / price)

    elif method == "equal":
        # Equal weight across n positions
        n_positions = kwargs.get("n_positions", 10)
        position_value = capital / n_positions
        return int(position_value / price)

    elif method == "volatility":
        # Inverse volatility sizing
        if volatility is None or volatility <= 0:
            raise ValueError("Volatility required for volatility-based sizing")
        target_vol = kwargs.get("target_volatility", 0.02)  # 2% daily vol
        position_value = capital * (target_vol / volatility)
        return int(position_value / price)

    elif method == "risk":
        # Risk-based sizing (Kelly-like)
        if stop_loss_pct is None or stop_loss_pct <= 0:
            raise ValueError("Stop loss percentage required for risk-based sizing")
        max_risk = capital * risk_per_trade
        risk_per_share = price * stop_loss_pct
        return int(max_risk / risk_per_share)

    else:
        raise ValueError(f"Invalid method: {method}")


def kelly_criterion(
    win_rate: float,
    avg_win: float,
    avg_loss: float,
    fraction: float = 0.5,
) -> float:
    """Calculate Kelly criterion position size.

    Args:
        win_rate: Probability of winning (0-1).
        avg_win: Average winning return.
        avg_loss: Average losing return (positive value).
        fraction: Fraction of Kelly to use (default 0.5 = half-Kelly).

    Returns:
        Optimal position size as fraction of capital.
    """
    if win_rate <= 0 or win_rate >= 1:
        return 0.0
    if avg_loss <= 0:
        return 0.0

    # Kelly formula: f* = W - (1-W)/R
    # where W = win rate, R = win/loss ratio
    win_loss_ratio = abs(avg_win / avg_loss) if avg_loss != 0 else 0
    kelly = win_rate - (1 - win_rate) / win_loss_ratio if win_loss_ratio > 0 else 0

    # Apply fractional Kelly and bound
    kelly = kelly * fraction
    return max(0, min(kelly, 1))


def risk_parity_weights(
    returns: pd.DataFrame,
    risk_budget: Mapping[str, float] | None = None,
) -> pd.Series:
    """Calculate risk parity portfolio weights.

    Each asset contributes equally to portfolio risk.

    Args:
        returns: DataFrame of asset returns (columns are assets).
        risk_budget: Optional risk budget per asset (defaults to equal).

    Returns:
        Series of portfolio weights.
    """
    # Calculate volatility for each asset
    vols = returns.std()

    # Inverse volatility weights (simple risk parity)
    inv_vol = 1 / vols
    weights = inv_vol / inv_vol.sum()

    # Apply risk budget if provided
    if risk_budget is not None:
        budget = pd.Series(risk_budget)
        budget = budget / budget.sum()  # Normalize
        weights = weights * budget
        weights = weights / weights.sum()  # Renormalize

    return weights


def correlation_matrix(
    returns: pd.DataFrame,
    method: str = "pearson",
) -> pd.DataFrame:
    """Calculate correlation matrix for returns.

    Args:
        returns: DataFrame of asset returns.
        method: Correlation method ('pearson', 'spearman', 'kendall').

    Returns:
        Correlation matrix as DataFrame.
    """
    return returns.corr(method=method)


def rolling_correlation(
    returns1: pd.Series,
    returns2: pd.Series,
    window: int = 60,
) -> pd.Series:
    """Calculate rolling correlation between two return series.

    Args:
        returns1: First return series.
        returns2: Second return series.
        window: Rolling window size.

    Returns:
        Series of rolling correlations.
    """
    return returns1.rolling(window=window).corr(returns2)


def diversification_ratio(
    weights: pd.Series,
    returns: pd.DataFrame,
) -> float:
    """Calculate portfolio diversification ratio.

    DR = weighted average volatility / portfolio volatility
    DR > 1 indicates diversification benefit.

    Args:
        weights: Portfolio weights.
        returns: DataFrame of asset returns.

    Returns:
        Diversification ratio.
    """
    # Individual asset volatilities
    vols = returns.std()

    # Weighted average volatility
    weighted_avg_vol = (weights * vols).sum()

    # Portfolio volatility
    cov_matrix = returns.cov()
    portfolio_var = weights @ cov_matrix @ weights
    portfolio_vol = np.sqrt(portfolio_var)

    if portfolio_vol == 0:
        return 1.0

    return weighted_avg_vol / portfolio_vol


def calculate_information_ratio(
    strategy_returns: pd.Series,
    benchmark_returns: pd.Series,
    periods_per_year: int = 252,
) -> float:
    """Calculate information ratio vs benchmark.

    IR = (strategy return - benchmark return) / tracking error

    Args:
        strategy_returns: Strategy return series.
        benchmark_returns: Benchmark return series.
        periods_per_year: Trading periods per year.

    Returns:
        Information ratio (annualized).
    """
    # Align series
    combined = pd.concat([strategy_returns, benchmark_returns], axis=1).dropna()
    if len(combined) < 2:
        return 0.0

    strat = combined.iloc[:, 0]
    bench = combined.iloc[:, 1]

    # Active returns
    active_returns = strat - bench

    # Tracking error
    tracking_error = active_returns.std() * np.sqrt(periods_per_year)

    if tracking_error == 0:
        return 0.0

    # Annualized active return
    active_return_annual = active_returns.mean() * periods_per_year

    return active_return_annual / tracking_error


def calculate_beta(
    strategy_returns: pd.Series,
    market_returns: pd.Series,
) -> float:
    """Calculate portfolio beta vs market.

    Args:
        strategy_returns: Strategy return series.
        market_returns: Market return series.

    Returns:
        Beta coefficient.
    """
    combined = pd.concat([strategy_returns, market_returns], axis=1).dropna()
    if len(combined) < 2:
        return 1.0

    strat = combined.iloc[:, 0]
    market = combined.iloc[:, 1]

    covariance = strat.cov(market)
    market_variance = market.var()

    if market_variance == 0:
        return 0.0

    return covariance / market_variance


def calculate_alpha(
    strategy_returns: pd.Series,
    market_returns: pd.Series,
    risk_free_rate: float = 0.0,
    periods_per_year: int = 252,
) -> float:
    """Calculate Jensen's alpha.

    Args:
        strategy_returns: Strategy return series.
        market_returns: Market return series.
        risk_free_rate: Annual risk-free rate.
        periods_per_year: Trading periods per year.

    Returns:
        Annualized alpha.
    """
    beta = calculate_beta(strategy_returns, market_returns)

    rf_per_period = risk_free_rate / periods_per_year

    # Alpha = strategy excess return - beta * market excess return
    strategy_excess = strategy_returns.mean() - rf_per_period
    market_excess = market_returns.mean() - rf_per_period

    alpha_per_period = strategy_excess - beta * market_excess

    return alpha_per_period * periods_per_year
