"""Weight allocation algorithms for strategy portfolio management.

This module provides various weight allocation methods including equal weight,
risk parity, momentum weighting, and custom weight calculations.
"""

from enum import Enum, auto
from typing import Dict, List, Optional, Protocol

import numpy as np
import pandas as pd


class WeightMethod(str, Enum):
    """Supported weight allocation methods."""

    EQUAL = "equal"
    RISK_PARITY = "risk_parity"
    MOMENTUM = "momentum"
    INVERSE_VOLATILITY = "inverse_volatility"
    CUSTOM = "custom"


class WeightCalculator(Protocol):
    """Protocol for weight calculation functions."""

    def __call__(
        self,
        strategies: List[str],
        returns_data: Optional[pd.DataFrame] = None,
        **kwargs: float,
    ) -> Dict[str, float]: ...


def calculate_equal_weights(
    strategies: List[str],
    returns_data: Optional[pd.DataFrame] = None,
    **kwargs: float,
) -> Dict[str, float]:
    """Calculate equal weights for all strategies.

    Args:
        strategies: List of strategy identifiers
        returns_data: Optional returns data (not used for equal weights)
        **kwargs: Additional parameters (not used)

    Returns:
        Dictionary mapping strategy names to equal weights
    """
    if not strategies:
        return {}

    weight = 1.0 / len(strategies)
    return {strategy: weight for strategy in strategies}


def calculate_risk_parity_weights(
    strategies: List[str],
    returns_data: Optional[pd.DataFrame] = None,
    lookback_period: int = 60,
    **kwargs: float,
) -> Dict[str, float]:
    """Calculate risk parity weights based on inverse volatility.

    Risk parity allocates weights inversely proportional to each strategy's
    risk (volatility), aiming for equal risk contribution from each strategy.

    Args:
        strategies: List of strategy identifiers
        returns_data: DataFrame with strategy returns (columns = strategies)
        lookback_period: Number of periods for volatility calculation
        **kwargs: Additional parameters

    Returns:
        Dictionary mapping strategy names to risk parity weights

    Raises:
        ValueError: If returns_data is not provided or insufficient data
    """
    if returns_data is None or returns_data.empty:
        # Fallback to equal weights if no data
        return calculate_equal_weights(strategies)

    # Filter to requested strategies
    available_strategies = [s for s in strategies if s in returns_data.columns]
    if not available_strategies:
        return calculate_equal_weights(strategies)

    # Calculate volatilities
    data = returns_data[available_strategies].dropna()
    if len(data) < lookback_period:
        # Use available data if less than lookback
        lookback_period = len(data)

    if lookback_period < 2:
        return calculate_equal_weights(available_strategies)

    # Calculate rolling volatility (standard deviation of returns)
    recent_returns = data.tail(lookback_period)
    volatilities = recent_returns.std()

    # Handle zero volatility case
    volatilities = volatilities.replace(0, np.nan)
    if volatilities.isna().all():
        return calculate_equal_weights(available_strategies)

    # Fill NaN volatilities with max volatility (lower weight)
    max_vol = volatilities.max()
    volatilities = volatilities.fillna(max_vol * 2)

    # Calculate inverse volatility weights
    inv_vol = 1.0 / volatilities
    weights = inv_vol / inv_vol.sum()

    result = dict(weights)

    # Add zero weights for strategies without data
    for strategy in strategies:
        if strategy not in result:
            result[strategy] = 0.0

    return result


def calculate_momentum_weights(
    strategies: List[str],
    returns_data: Optional[pd.DataFrame] = None,
    lookback_period: int = 90,
    momentum_decay: float = 1.0,
    **kwargs: float,
) -> Dict[str, float]:
    """Calculate momentum-based weights.

    Allocates higher weights to strategies with better recent performance.

    Args:
        strategies: List of strategy identifiers
        returns_data: DataFrame with strategy returns (columns = strategies)
        lookback_period: Number of periods for momentum calculation
        momentum_decay: Decay factor for older returns (1.0 = no decay)
        **kwargs: Additional parameters

    Returns:
        Dictionary mapping strategy names to momentum weights
    """
    if returns_data is None or returns_data.empty:
        return calculate_equal_weights(strategies)

    # Filter to requested strategies
    available_strategies = [s for s in strategies if s in returns_data.columns]
    if not available_strategies:
        return calculate_equal_weights(strategies)

    # Calculate momentum (cumulative return)
    data = returns_data[available_strategies].dropna()
    if len(data) < 2:
        return calculate_equal_weights(available_strategies)

    use_period = min(lookback_period, len(data))
    recent_returns = data.tail(use_period)

    # Apply decay weights if specified
    if momentum_decay != 1.0:
        decay_weights = np.power(
            momentum_decay, np.arange(len(recent_returns))[::-1]
        )
        decay_weights = decay_weights / decay_weights.sum()

        momentum = (recent_returns * decay_weights[:, np.newaxis]).sum()
    else:
        # Simple cumulative return
        momentum = (1 + recent_returns).prod() - 1

    # Handle negative momentum - shift to positive
    min_momentum = momentum.min()
    if min_momentum < 0:
        momentum = momentum - min_momentum + 0.001

    # Calculate weights proportional to momentum
    if momentum.sum() > 0:
        weights = momentum / momentum.sum()
    else:
        return calculate_equal_weights(available_strategies)

    result = dict(weights)

    # Add zero weights for strategies without data
    for strategy in strategies:
        if strategy not in result:
            result[strategy] = 0.0

    return result


def calculate_inverse_volatility_weights(
    strategies: List[str],
    returns_data: Optional[pd.DataFrame] = None,
    lookback_period: int = 60,
    **kwargs: float,
) -> Dict[str, float]:
    """Calculate weights based on inverse volatility.

    Similar to risk parity but simpler - just uses inverse of volatility.

    Args:
        strategies: List of strategy identifiers
        returns_data: DataFrame with strategy returns (columns = strategies)
        lookback_period: Number of periods for volatility calculation
        **kwargs: Additional parameters

    Returns:
        Dictionary mapping strategy names to inverse volatility weights
    """
    if returns_data is None or returns_data.empty:
        return calculate_equal_weights(strategies)

    available_strategies = [s for s in strategies if s in returns_data.columns]
    if not available_strategies:
        return calculate_equal_weights(strategies)

    data = returns_data[available_strategies].dropna()
    if len(data) < 2:
        return calculate_equal_weights(available_strategies)

    use_period = min(lookback_period, len(data))
    recent_returns = data.tail(use_period)

    # Calculate volatility
    volatilities = recent_returns.std()

    # Handle edge cases
    volatilities = volatilities.replace(0, np.nan).fillna(volatilities.max() * 2)

    # Calculate inverse volatility weights
    inv_vol = 1.0 / volatilities
    weights = inv_vol / inv_vol.sum()

    result = dict(weights)

    # Add zero weights for missing strategies
    for strategy in strategies:
        if strategy not in result:
            result[strategy] = 0.0

    return result


def validate_weights(weights: Dict[str, float], tolerance: float = 0.01) -> bool:
    """Validate that weights sum to approximately 1.0.

    Args:
        weights: Dictionary of strategy weights
        tolerance: Allowed deviation from 1.0

    Returns:
        True if weights are valid
    """
    if not weights:
        return True

    total = sum(weights.values())
    return abs(total - 1.0) <= tolerance


def normalize_weights(weights: Dict[str, float]) -> Dict[str, float]:
    """Normalize weights to sum to 1.0.

    Args:
        weights: Dictionary of strategy weights

    Returns:
        Normalized weights dictionary
    """
    if not weights:
        return {}

    total = sum(weights.values())
    if total == 0:
        return calculate_equal_weights(list(weights.keys()))

    return {k: v / total for k, v in weights.items()}


def apply_weight_constraints(
    weights: Dict[str, float],
    min_weight: float = 0.0,
    max_weight: float = 1.0,
) -> Dict[str, float]:
    """Apply minimum and maximum weight constraints with iterative redistribution.

    This function applies constraints by:
    1. Clipping weights to [min, max]
    2. Iteratively redistributing excess/deficit to unconstrained strategies
    3. Stopping when all constraints are satisfied or no more redistribution possible

    Args:
        weights: Dictionary of strategy weights
        min_weight: Minimum weight for each strategy
        max_weight: Maximum weight for each strategy

    Returns:
        Constrained weights (may not sum to exactly 1.0 if constraints are infeasible)
    """
    if not weights:
        return {}

    if len(weights) == 1:
        return {list(weights.keys())[0]: 1.0}

    # Start with a copy and normalize
    result = normalize_weights(weights)
    max_iterations = len(weights) * 3

    for _ in range(max_iterations):
        # Clip weights to constraints
        clipped = {k: max(min_weight, min(max_weight, v)) for k, v in result.items()}

        # Check for convergence
        if all(abs(clipped[k] - result[k]) < 1e-10 for k in result):
            return clipped

        # Calculate how much weight was clipped
        clipped_sum = sum(clipped.values())
        excess = 1.0 - clipped_sum

        # Find strategies that aren't at bounds (can receive redistributed weight)
        unconstrained = []
        for k, v in clipped.items():
            # Can receive more weight if not at max and excess > 0
            # Can give up weight if not at min and excess < 0
            if excess > 0 and v < max_weight - 1e-10:
                unconstrained.append(k)
            elif excess < 0 and v > min_weight + 1e-10:
                unconstrained.append(k)

        if not unconstrained or abs(excess) < 1e-10:
            # No place to redistribute or no excess
            return clipped

        # Redistribute excess proportionally among unconstrained strategies
        unconstrained_sum = sum(clipped[k] for k in unconstrained)
        if unconstrained_sum > 0:
            for k in unconstrained:
                proportion = clipped[k] / unconstrained_sum
                clipped[k] += excess * proportion
        else:
            # Distribute equally if sum is 0
            for k in unconstrained:
                clipped[k] += excess / len(unconstrained)

        result = clipped

    # Final clip after iterations
    return {k: max(min_weight, min(max_weight, v)) for k, v in result.items()}
