"""Rebalancing module for strategy portfolio management.

This module provides rebalancing logic including periodic rebalancing,
threshold-based triggers, and transaction cost calculations.
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum, auto
from typing import Callable, Dict, List, Optional, Protocol

import numpy as np


class RebalanceTrigger(str, Enum):
    """Types of rebalancing triggers."""

    PERIODIC = "periodic"
    THRESHOLD = "threshold"
    CALENDAR = "calendar"
    MANUAL = "manual"


@dataclass
class RebalanceResult:
    """Result of a rebalancing operation.

    Attributes:
        timestamp: When rebalancing occurred
        old_weights: Weights before rebalancing
        new_weights: Target weights after rebalancing
        actual_weights: Actual achieved weights (accounting for costs)
        trades_executed: Number of trades executed
        transaction_costs: Total transaction costs
        drift_before: Portfolio drift before rebalancing
        trigger: What triggered the rebalance
    """

    timestamp: datetime
    old_weights: Dict[str, float]
    new_weights: Dict[str, float]
    actual_weights: Dict[str, float] = field(default_factory=dict)
    trades_executed: int = 0
    transaction_costs: float = 0.0
    drift_before: float = 0.0
    trigger: RebalanceTrigger = RebalanceTrigger.MANUAL

    @property
    def weight_changes(self) -> Dict[str, float]:
        """Calculate weight changes from rebalancing."""
        all_strategies = set(self.old_weights.keys()) | set(self.new_weights.keys())
        return {
            s: self.new_weights.get(s, 0.0) - self.old_weights.get(s, 0.0)
            for s in all_strategies
        }

    @property
    def total_turnover(self) -> float:
        """Calculate total portfolio turnover."""
        return sum(abs(change) for change in self.weight_changes.values()) / 2.0


@dataclass
class TransactionCostModel:
    """Model for calculating transaction costs.

    Attributes:
        fixed_cost: Fixed cost per trade
        percentage_cost: Percentage of trade value
        market_impact_factor: Factor for market impact cost
        min_cost: Minimum cost per trade
        max_cost: Maximum cost per trade (None = no limit)
    """

    fixed_cost: float = 0.0
    percentage_cost: float = 0.001
    market_impact_factor: float = 0.0
    min_cost: float = 0.0
    max_cost: Optional[float] = None

    def calculate_cost(
        self,
        trade_value: float,
        strategy_volatility: float = 0.0,
    ) -> float:
        """Calculate transaction cost for a trade.

        Args:
            trade_value: Value of the trade
            strategy_volatility: Volatility of the strategy (for impact)

        Returns:
            Total transaction cost
        """
        # Fixed + percentage cost
        cost = self.fixed_cost + abs(trade_value) * self.percentage_cost

        # Market impact cost
        if self.market_impact_factor > 0 and strategy_volatility > 0:
            impact = abs(trade_value) * strategy_volatility * self.market_impact_factor
            cost += impact

        # Apply min/max constraints
        cost = max(cost, self.min_cost)
        if self.max_cost is not None:
            cost = min(cost, self.max_cost)

        return cost


class Rebalancer:
    """Portfolio rebalancer with multiple trigger mechanisms.

    This class manages portfolio rebalancing based on various triggers
    including periodic intervals, drift thresholds, and manual requests.

    Args:
        trigger_type: Primary rebalancing trigger type
        rebalance_frequency: Days between periodic rebalances
        drift_threshold: Maximum allowed drift before rebalancing
        cost_model: Transaction cost model
        min_rebalance_interval: Minimum days between rebalances
    """

    def __init__(
        self,
        trigger_type: RebalanceTrigger = RebalanceTrigger.PERIODIC,
        rebalance_frequency: int = 30,
        drift_threshold: float = 0.05,
        cost_model: Optional[TransactionCostModel] = None,
        min_rebalance_interval: int = 1,
    ):
        self.trigger_type = trigger_type
        self.rebalance_frequency = rebalance_frequency
        self.drift_threshold = drift_threshold
        self.cost_model = cost_model or TransactionCostModel()
        self.min_rebalance_interval = min_rebalance_interval

        # State tracking
        self._last_rebalance_time: Optional[datetime] = None
        self._target_weights: Dict[str, float] = {}
        self._current_weights: Dict[str, float] = {}
        self._rebalance_history: List[RebalanceResult] = []

    def check_rebalance_needed(
        self,
        current_weights: Dict[str, float],
        target_weights: Dict[str, float],
        current_time: Optional[datetime] = None,
    ) -> tuple[bool, RebalanceTrigger]:
        """Check if rebalancing is needed.

        Args:
            current_weights: Current portfolio weights
            target_weights: Target portfolio weights
            current_time: Current timestamp (default: now)

        Returns:
            Tuple of (rebalance_needed, trigger_type)
        """
        if current_time is None:
            current_time = datetime.now()

        # Check minimum interval
        if self._last_rebalance_time is not None:
            days_since_last = (current_time - self._last_rebalance_time).days
            if days_since_last < self.min_rebalance_interval:
                return False, RebalanceTrigger.MANUAL

        # Check periodic trigger
        if self.trigger_type == RebalanceTrigger.PERIODIC:
            if self._last_rebalance_time is None:
                return True, RebalanceTrigger.PERIODIC

            days_since_last = (current_time - self._last_rebalance_time).days
            if days_since_last >= self.rebalance_frequency:
                return True, RebalanceTrigger.PERIODIC

        # Check threshold trigger
        if self.trigger_type in (RebalanceTrigger.THRESHOLD, RebalanceTrigger.PERIODIC):
            drift = self.calculate_drift(current_weights, target_weights)
            if drift > self.drift_threshold:
                return True, RebalanceTrigger.THRESHOLD

        return False, RebalanceTrigger.MANUAL

    def calculate_drift(
        self,
        current_weights: Dict[str, float],
        target_weights: Dict[str, float],
    ) -> float:
        """Calculate portfolio drift from target.

        Args:
            current_weights: Current portfolio weights
            target_weights: Target portfolio weights

        Returns:
            Maximum absolute deviation (drift)
        """
        all_strategies = set(current_weights.keys()) | set(target_weights.keys())

        max_drift = 0.0
        for strategy in all_strategies:
            current = current_weights.get(strategy, 0.0)
            target = target_weights.get(strategy, 0.0)
            drift = abs(current - target)
            max_drift = max(max_drift, drift)

        return max_drift

    def rebalance(
        self,
        current_weights: Dict[str, float],
        target_weights: Dict[str, float],
        portfolio_value: float,
        current_time: Optional[datetime] = None,
        strategy_volatilities: Optional[Dict[str, float]] = None,
        force: bool = False,
    ) -> Optional[RebalanceResult]:
        """Execute portfolio rebalancing.

        Args:
            current_weights: Current portfolio weights
            target_weights: Target portfolio weights
            portfolio_value: Total portfolio value
            current_time: Current timestamp
            strategy_volatilities: Optional volatility per strategy
            force: Force rebalance even if not triggered

        Returns:
            RebalanceResult if rebalancing occurred, None otherwise
        """
        if current_time is None:
            current_time = datetime.now()

        # Check if rebalance is needed (unless forced)
        if not force:
            needed, trigger = self.check_rebalance_needed(
                current_weights, target_weights, current_time
            )
            if not needed:
                return None
        else:
            trigger = RebalanceTrigger.MANUAL

        # Calculate drift
        drift = self.calculate_drift(current_weights, target_weights)

        # Calculate trades and costs
        trades_executed = 0
        total_costs = 0.0
        volatilities = strategy_volatilities or {}

        all_strategies = set(current_weights.keys()) | set(target_weights.keys())

        for strategy in all_strategies:
            current = current_weights.get(strategy, 0.0)
            target = target_weights.get(strategy, 0.0)
            change = target - current

            if abs(change) > 0.0001:  # Minimum trade threshold
                trades_executed += 1
                trade_value = abs(change) * portfolio_value
                vol = volatilities.get(strategy, 0.0)
                total_costs += self.cost_model.calculate_cost(trade_value, vol)

        # Create result (actual weights account for costs)
        result = RebalanceResult(
            timestamp=current_time,
            old_weights=current_weights.copy(),
            new_weights=target_weights.copy(),
            actual_weights=target_weights.copy(),  # Simplified - in reality costs reduce weights
            trades_executed=trades_executed,
            transaction_costs=total_costs,
            drift_before=drift,
            trigger=trigger,
        )

        # Update state
        self._last_rebalance_time = current_time
        self._current_weights = target_weights.copy()
        self._target_weights = target_weights.copy()
        self._rebalance_history.append(result)

        return result

    def get_rebalance_history(
        self,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
    ) -> List[RebalanceResult]:
        """Get rebalancing history within time range.

        Args:
            start_time: Start of time range (inclusive)
            end_time: End of time range (inclusive)

        Returns:
            List of rebalance results
        """
        history = self._rebalance_history

        if start_time:
            history = [h for h in history if h.timestamp >= start_time]
        if end_time:
            history = [h for h in history if h.timestamp <= end_time]

        return history

    def get_total_transaction_costs(
        self,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
    ) -> float:
        """Calculate total transaction costs over period.

        Args:
            start_time: Start of time range
            end_time: End of time range

        Returns:
            Total transaction costs
        """
        history = self.get_rebalance_history(start_time, end_time)
        return sum(h.transaction_costs for h in history)

    def get_average_turnover(
        self,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
    ) -> float:
        """Calculate average portfolio turnover.

        Args:
            start_time: Start of time range
            end_time: End of time range

        Returns:
            Average turnover per rebalance
        """
        history = self.get_rebalance_history(start_time, end_time)
        if not history:
            return 0.0
        return float(np.mean([h.total_turnover for h in history]))

    def reset(self) -> None:
        """Reset rebalancer state."""
        self._last_rebalance_time = None
        self._target_weights = {}
        self._current_weights = {}
        self._rebalance_history = []

    def set_target_weights(self, weights: Dict[str, float]) -> None:
        """Set target weights for the portfolio.

        Args:
            weights: Dictionary of strategy weights
        """
        self._target_weights = weights.copy()

    def update_config(
        self,
        trigger_type: Optional[RebalanceTrigger] = None,
        rebalance_frequency: Optional[int] = None,
        drift_threshold: Optional[float] = None,
    ) -> None:
        """Update rebalancer configuration.

        Args:
            trigger_type: New trigger type
            rebalance_frequency: New frequency in days
            drift_threshold: New drift threshold
        """
        if trigger_type is not None:
            self.trigger_type = trigger_type
        if rebalance_frequency is not None:
            self.rebalance_frequency = rebalance_frequency
        if drift_threshold is not None:
            self.drift_threshold = drift_threshold
