"""Fitness evaluation module for OpenClaw evolution system.

This module provides fitness evaluation functions for trading strategies,
including multi-objective metrics like profit, risk, and drawdown.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple
from datetime import datetime

import numpy as np
import pandas as pd

from openclaw.backtest.engine import BacktestResult
from openclaw.utils.logging import get_logger


@dataclass
class FitnessMetrics:
    """Comprehensive fitness metrics for a trading strategy.

    Attributes:
        total_return: Total return percentage
        sharpe_ratio: Sharpe ratio
        max_drawdown: Maximum drawdown percentage
        win_rate: Win rate percentage
        profit_factor: Profit factor
        calmar_ratio: Calmar ratio
        volatility: Annualized volatility
        num_trades: Number of trades
        avg_trade_return: Average return per trade
        fitness_score: Combined fitness score
    """

    total_return: float = 0.0
    sharpe_ratio: float = 0.0
    max_drawdown: float = 0.0
    win_rate: float = 0.0
    profit_factor: float = 0.0
    calmar_ratio: float = 0.0
    volatility: float = 0.0
    num_trades: int = 0
    avg_trade_return: float = 0.0
    fitness_score: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        """Convert metrics to dictionary."""
        return {
            "total_return": self.total_return,
            "sharpe_ratio": self.sharpe_ratio,
            "max_drawdown": self.max_drawdown,
            "win_rate": self.win_rate,
            "profit_factor": self.profit_factor,
            "calmar_ratio": self.calmar_ratio,
            "volatility": self.volatility,
            "num_trades": self.num_trades,
            "avg_trade_return": self.avg_trade_return,
            "fitness_score": self.fitness_score,
        }


class FitnessFunction(Protocol):
    """Protocol for fitness evaluation functions."""

    def __call__(
        self,
        backtest_result: Optional[BacktestResult] = None,
        metrics: Optional[FitnessMetrics] = None,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness score.

        Args:
            backtest_result: Backtest results
            metrics: Pre-calculated metrics
            **kwargs: Additional parameters

        Returns:
            Fitness score (higher is better)
        """
        ...


class FitnessEvaluator:
    """Evaluates fitness of trading strategies.

    This class provides various fitness functions and metrics calculation
    for evaluating trading strategy performance.

    Args:
        risk_free_rate: Annual risk-free rate (default: 0.02 for 2%)
        min_trades: Minimum trades required for valid fitness
    """

    def __init__(
        self,
        risk_free_rate: float = 0.02,
        min_trades: int = 10,
    ):
        self.risk_free_rate = risk_free_rate
        self.min_trades = min_trades
        self.logger = get_logger("evolution.fitness")

    def calculate_metrics(self, backtest_result: BacktestResult) -> FitnessMetrics:
        """Calculate comprehensive fitness metrics from backtest results.

        Args:
            backtest_result: Backtest results

        Returns:
            FitnessMetrics object
        """
        if backtest_result.total_trades < self.min_trades:
            return FitnessMetrics()

        metrics = FitnessMetrics(
            total_return=backtest_result.total_return,
            sharpe_ratio=backtest_result.sharpe_ratio,
            max_drawdown=backtest_result.max_drawdown,
            win_rate=backtest_result.win_rate,
            profit_factor=backtest_result.profit_factor,
            calmar_ratio=backtest_result.calmar_ratio,
            volatility=backtest_result.volatility,
            num_trades=backtest_result.total_trades,
            avg_trade_return=(
                backtest_result.total_return / backtest_result.total_trades
                if backtest_result.total_trades > 0 else 0.0
            ),
        )

        return metrics

    def calculate_fitness_sharpe(
        self,
        backtest_result: Optional[BacktestResult] = None,
        metrics: Optional[FitnessMetrics] = None,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness based on Sharpe ratio.

        Args:
            backtest_result: Backtest results
            metrics: Pre-calculated metrics

        Returns:
            Fitness score
        """
        if metrics is None and backtest_result is not None:
            metrics = self.calculate_metrics(backtest_result)

        if metrics is None or metrics.num_trades < self.min_trades:
            return -1.0

        # Penalize high drawdown
        drawdown_penalty = max(0, metrics.max_drawdown - 20) * 0.05

        # Base fitness on Sharpe ratio
        fitness = metrics.sharpe_ratio - drawdown_penalty

        return fitness

    def calculate_fitness_profit_risk(
        self,
        backtest_result: Optional[BacktestResult] = None,
        metrics: Optional[FitnessMetrics] = None,
        risk_weight: float = 0.5,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness balancing profit and risk.

        Args:
            backtest_result: Backtest results
            metrics: Pre-calculated metrics
            risk_weight: Weight for risk component (0-1)

        Returns:
            Fitness score
        """
        if metrics is None and backtest_result is not None:
            metrics = self.calculate_metrics(backtest_result)

        if metrics is None or metrics.num_trades < self.min_trades:
            return -1.0

        # Normalize return to 0-1 range (assuming -100% to +100%)
        return_score = (metrics.total_return + 100) / 200
        return_score = max(0, min(1, return_score))

        # Risk score (inverse of drawdown, normalized)
        risk_score = 1 - (metrics.max_drawdown / 100)
        risk_score = max(0, min(1, risk_score))

        # Combined fitness
        profit_weight = 1 - risk_weight
        fitness = profit_weight * return_score + risk_weight * risk_score

        # Bonus for good Sharpe ratio
        if metrics.sharpe_ratio > 1.0:
            fitness *= 1.2

        return fitness

    def calculate_fitness_multi_objective(
        self,
        backtest_result: Optional[BacktestResult] = None,
        metrics: Optional[FitnessMetrics] = None,
        objectives: Optional[List[str]] = None,
        weights: Optional[List[float]] = None,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness using multiple weighted objectives.

        Args:
            backtest_result: Backtest results
            metrics: Pre-calculated metrics
            objectives: List of objective names
            weights: List of objective weights

        Returns:
            Fitness score
        """
        if metrics is None and backtest_result is not None:
            metrics = self.calculate_metrics(backtest_result)

        if metrics is None or metrics.num_trades < self.min_trades:
            return -1.0

        if objectives is None:
            objectives = ["total_return", "sharpe_ratio", "win_rate"]

        if weights is None:
            weights = [0.4, 0.4, 0.2]

        if len(objectives) != len(weights):
            raise ValueError("Number of objectives must match number of weights")

        # Normalize weights
        total_weight = sum(weights)
        weights = [w / total_weight for w in weights]

        scores = []
        for obj in objectives:
            value = getattr(metrics, obj, 0.0)

            # Normalize different metrics to similar scales
            if obj == "total_return":
                score = max(-1, min(1, value / 100))  # -100% to +100%
            elif obj == "sharpe_ratio":
                score = max(-1, min(2, value / 2))  # -2 to +4
            elif obj == "max_drawdown":
                score = max(-1, min(0, -value / 50))  # 0 to -50%
            elif obj == "win_rate":
                score = (value - 50) / 50  # 50% is neutral, 100% is +1
            elif obj == "profit_factor":
                score = min(1, (value - 1) / 2)  # 1 is neutral, 3 is +1
            elif obj == "calmar_ratio":
                score = min(1, value / 3)  # 0 to 3
            elif obj == "volatility":
                score = max(-1, 0.3 - value / 100)  # Lower is better
            else:
                score = value

            scores.append(score)

        # Weighted sum
        fitness = sum(s * w for s, w in zip(scores, weights))

        # Ensure minimum trades
        trade_penalty = max(0, self.min_trades - metrics.num_trades) / self.min_trades
        fitness -= trade_penalty * 0.5

        return fitness

    def calculate_fitness_sortino(
        self,
        backtest_result: Optional[BacktestResult] = None,
        returns: Optional[np.ndarray] = None,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness based on Sortino ratio.

        Args:
            backtest_result: Backtest results
            returns: Array of returns

        Returns:
            Fitness score
        """
        if backtest_result is not None and backtest_result.total_trades < self.min_trades:
            return -1.0

        if returns is None and backtest_result is not None:
            # Calculate returns from equity curve
            equity = np.array(backtest_result.equity_curve)
            returns = np.diff(equity) / equity[:-1]

        if returns is None or len(returns) < 2:
            return -1.0

        # Calculate Sortino ratio
        excess_returns = returns - self.risk_free_rate / 252  # Daily
        downside_returns = excess_returns[excess_returns < 0]

        if len(downside_returns) == 0 or downside_returns.std() == 0:
            return excess_returns.mean() * 252  # Annualized

        downside_deviation = downside_returns.std() * np.sqrt(252)
        expected_return = excess_returns.mean() * 252

        sortino = expected_return / downside_deviation

        return sortino

    def calculate_fitness_robustness(
        self,
        backtest_results: List[BacktestResult],
        **kwargs: Any,
    ) -> float:
        """Calculate fitness based on robustness across multiple periods.

        Args:
            backtest_results: List of backtest results for different periods

        Returns:
            Fitness score
        """
        if len(backtest_results) < 2:
            return -1.0

        # Calculate fitness for each period
        period_fitnesses = [
            self.calculate_fitness_sharpe(result) for result in backtest_results
        ]

        # Filter out invalid results
        valid_fitnesses = [f for f in period_fitnesses if f > -0.5]

        if len(valid_fitnesses) < len(backtest_results) * 0.5:
            return -1.0

        # Consistency score: how consistent is performance across periods
        mean_fitness = np.mean(valid_fitnesses)
        std_fitness = np.std(valid_fitnesses)
        consistency = 1 / (1 + std_fitness)

        # Robust fitness: weighted combination of mean and consistency
        robust_fitness = 0.7 * mean_fitness + 0.3 * consistency

        return robust_fitness

    def calculate_fitness_adaptive(
        self,
        backtest_result: Optional[BacktestResult] = None,
        metrics: Optional[FitnessMetrics] = None,
        generation: int = 0,
        max_generations: int = 500,
        **kwargs: Any,
    ) -> float:
        """Calculate fitness with adaptive weighting based on evolution progress.

        Early generations favor exploration (riskier strategies),
        later generations favor exploitation (stable strategies).

        Args:
            backtest_result: Backtest results
            metrics: Pre-calculated metrics
            generation: Current generation
            max_generations: Maximum generations

        Returns:
            Fitness score
        """
        if metrics is None and backtest_result is not None:
            metrics = self.calculate_metrics(backtest_result)

        if metrics is None or metrics.num_trades < self.min_trades:
            return -1.0

        # Adaptive risk weight
        progress = generation / max_generations if max_generations > 0 else 1.0
        risk_weight = 0.3 + 0.4 * progress  # 0.3 to 0.7

        return self.calculate_fitness_profit_risk(
            backtest_result=backtest_result,
            metrics=metrics,
            risk_weight=risk_weight,
        )

    def create_fitness_function(
        self,
        function_type: str = "sharpe",
        **kwargs: Any,
    ) -> Callable[..., float]:
        """Create a fitness function of the specified type.

        Args:
            function_type: Type of fitness function
            **kwargs: Additional parameters for the fitness function

        Returns:
            Fitness function
        """
        fitness_functions: Dict[str, Callable[..., float]] = {
            "sharpe": self.calculate_fitness_sharpe,
            "profit_risk": self.calculate_fitness_profit_risk,
            "multi_objective": self.calculate_fitness_multi_objective,
            "sortino": self.calculate_fitness_sortino,
            "adaptive": self.calculate_fitness_adaptive,
        }

        if function_type not in fitness_functions:
            raise ValueError(f"Unknown fitness function type: {function_type}")

        base_func = fitness_functions[function_type]

        def fitness_wrapper(
            backtest_result: Optional[BacktestResult] = None,
            metrics: Optional[FitnessMetrics] = None,
            **extra_kwargs: Any,
        ) -> float:
            merged_kwargs = {**kwargs, **extra_kwargs}
            return base_func(backtest_result, metrics, **merged_kwargs)

        return fitness_wrapper

    def evaluate_population(
        self,
        backtest_results: List[BacktestResult],
        fitness_func: Optional[Callable[..., float]] = None,
    ) -> List[float]:
        """Evaluate fitness for a population of backtest results.

        Args:
            backtest_results: List of backtest results
            fitness_func: Optional custom fitness function

        Returns:
            List of fitness scores
        """
        if fitness_func is None:
            fitness_func = self.calculate_fitness_sharpe

        fitness_scores = []
        for result in backtest_results:
            try:
                score = fitness_func(backtest_result=result)
                fitness_scores.append(score)
            except Exception as e:
                self.logger.warning(f"Fitness evaluation failed: {e}")
                fitness_scores.append(-1.0)

        return fitness_scores

    def get_convergence_metrics(
        self,
        fitness_history: List[float],
        window: int = 20,
    ) -> Dict[str, float]:
        """Calculate convergence metrics from fitness history.

        Args:
            fitness_history: List of best fitness values over generations
            window: Window size for calculating trend

        Returns:
            Dictionary of convergence metrics
        """
        if len(fitness_history) < window:
            return {
                "converged": False,
                "improvement_rate": 0.0,
                "stability": 0.0,
            }

        recent = fitness_history[-window:]
        previous = fitness_history[-(window * 2) : -window]

        if len(previous) < window:
            previous = fitness_history[:window]

        recent_mean = np.mean(recent)
        previous_mean = np.mean(previous)

        improvement_rate = (recent_mean - previous_mean) / (abs(previous_mean) + 1e-10)
        stability = 1 - (np.std(recent) / (abs(recent_mean) + 1e-10))

        # Converged if low improvement and high stability
        converged = improvement_rate < 0.01 and stability > 0.95

        return {
            "converged": converged,
            "improvement_rate": float(improvement_rate),
            "stability": float(stability),
            "recent_mean": float(recent_mean),
        }
