"""Base optimizer implementation for OpenClaw Trading.

This module provides the base classes and parameter space definitions
for strategy optimization.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import numpy as np

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


class ParameterType(str, Enum):
    """Types of parameters in the parameter space."""

    CONTINUOUS = "continuous"
    DISCRETE = "discrete"
    CATEGORICAL = "categorical"
    INTEGER = "integer"


@dataclass
class ParameterRange:
    """Defines a range for a parameter.

    Attributes:
        name: Parameter name
        param_type: Type of parameter
        bounds: (min, max) for continuous/integer, or list of values for discrete/categorical
        step: Step size for discrete parameters (optional)
        distribution: Distribution for random sampling ("uniform", "log_uniform")
    """

    name: str
    param_type: ParameterType
    bounds: tuple[float, float] | list[Any]
    step: float | None = None
    distribution: str = "uniform"

    def __post_init__(self):
        """Validate parameter range."""
        if self.param_type in (ParameterType.CONTINUOUS, ParameterType.INTEGER):
            if not isinstance(self.bounds, tuple) or len(self.bounds) != 2:
                raise ValueError(f"{self.param_type.value} parameters need (min, max) bounds")
            if self.bounds[0] >= self.bounds[1]:
                raise ValueError(f"Invalid bounds: {self.bounds[0]} >= {self.bounds[1]}")
        elif not isinstance(self.bounds, list) or len(self.bounds) == 0:
            raise ValueError(f"{self.param_type.value} parameters need a list of values")

        if self.distribution not in ("uniform", "log_uniform"):
            raise ValueError(f"Unknown distribution: {self.distribution}")


class ParameterSpace:
    """Defines the parameter space for optimization.

    This class manages parameter ranges and generates parameter combinations
    for different optimization algorithms.
    """

    def __init__(self):
        """Initialize empty parameter space."""
        self._parameters: dict[str, ParameterRange] = {}
        self.logger = get_logger("optimizer.parameter_space")

    def add_continuous(
        self,
        name: str,
        low: float,
        high: float,
        distribution: str = "uniform",
    ) -> ParameterSpace:
        """Add a continuous parameter.

        Args:
            name: Parameter name
            low: Lower bound
            high: Upper bound
            distribution: Sampling distribution

        Returns:
            Self for method chaining
        """
        self._parameters[name] = ParameterRange(
            name=name,
            param_type=ParameterType.CONTINUOUS,
            bounds=(low, high),
            distribution=distribution,
        )
        return self

    def add_integer(
        self,
        name: str,
        low: int,
        high: int,
        distribution: str = "uniform",
    ) -> ParameterSpace:
        """Add an integer parameter.

        Args:
            name: Parameter name
            low: Lower bound
            high: Upper bound
            distribution: Sampling distribution

        Returns:
            Self for method chaining
        """
        self._parameters[name] = ParameterRange(
            name=name,
            param_type=ParameterType.INTEGER,
            bounds=(low, high),
            distribution=distribution,
        )
        return self

    def add_discrete(
        self,
        name: str,
        values: list[Any],
        step: float | None = None,
    ) -> ParameterSpace:
        """Add a discrete parameter.

        Args:
            name: Parameter name
            values: List of possible values
            step: Step size (optional)

        Returns:
            Self for method chaining
        """
        self._parameters[name] = ParameterRange(
            name=name,
            param_type=ParameterType.DISCRETE,
            bounds=values,
            step=step,
        )
        return self

    def add_categorical(self, name: str, choices: list[Any]) -> ParameterSpace:
        """Add a categorical parameter.

        Args:
            name: Parameter name
            choices: List of possible choices

        Returns:
            Self for method chaining
        """
        self._parameters[name] = ParameterRange(
            name=name,
            param_type=ParameterType.CATEGORICAL,
            bounds=choices,
        )
        return self

    def get_parameter(self, name: str) -> ParameterRange:
        """Get a parameter definition by name."""
        if name not in self._parameters:
            raise KeyError(f"Parameter '{name}' not found")
        return self._parameters[name]

    def sample_random(self) -> dict[str, Any]:
        """Sample random parameters from the space.

        Returns:
            Dictionary of parameter names to sampled values
        """
        params = {}
        for name, param_range in self._parameters.items():
            if param_range.param_type == ParameterType.CONTINUOUS:
                low, high = param_range.bounds
                if param_range.distribution == "log_uniform":
                    params[name] = np.exp(np.random.uniform(np.log(low), np.log(high)))
                else:
                    params[name] = np.random.uniform(low, high)
            elif param_range.param_type == ParameterType.INTEGER:
                low, high = param_range.bounds
                params[name] = np.random.randint(low, high + 1)
            elif param_range.param_type in (ParameterType.DISCRETE, ParameterType.CATEGORICAL):
                params[name] = np.random.choice(param_range.bounds)
        return params

    def get_grid_points(self, n_points: int | None = None) -> list[dict[str, Any]]:
        """Generate grid points for grid search.

        Args:
            n_points: Number of points per continuous dimension (default: auto)

        Returns:
            List of parameter dictionaries
        """
        grid_values = {}

        for name, param_range in self._parameters.items():
            if param_range.param_type == ParameterType.CONTINUOUS:
                low, high = param_range.bounds
                n = n_points or 5
                grid_values[name] = np.linspace(low, high, n).tolist()
            elif param_range.param_type == ParameterType.INTEGER:
                low, high = param_range.bounds
                grid_values[name] = list(range(low, high + 1))
            elif param_range.param_type in (ParameterType.DISCRETE, ParameterType.CATEGORICAL):
                grid_values[name] = list(param_range.bounds)

        # Generate all combinations
        from itertools import product

        keys = list(grid_values.keys())
        values = [grid_values[k] for k in keys]
        combinations = list(product(*values))

        return [dict(zip(keys, combo, strict=False)) for combo in combinations]

    def __len__(self) -> int:
        """Return number of parameters."""
        return len(self._parameters)

    def __contains__(self, name: str) -> bool:
        """Check if parameter exists."""
        return name in self._parameters


class OptimizationObjective(str, Enum):
    """Optimization objectives."""

    MAXIMIZE_RETURN = "maximize_return"
    MAXIMIZE_SHARPE = "maximize_sharpe"
    MAXIMIZE_CALMAR = "maximize_calmar"
    MINIMIZE_DRAWDOWN = "minimize_drawdown"
    MAXIMIZE_WIN_RATE = "maximize_win_rate"
    CUSTOM = "custom"


@dataclass
class OptimizerConfig:
    """Configuration for strategy optimization.

    Attributes:
        objective: Optimization objective
        max_iterations: Maximum number of iterations
        n_jobs: Number of parallel jobs (-1 for all cores)
        early_stopping: Whether to enable early stopping
        early_stopping_patience: Patience for early stopping
        early_stopping_min_delta: Minimum improvement for early stopping
        validation_split: Fraction of data to use for validation (0-1)
        random_state: Random seed for reproducibility
        custom_scorer: Custom scoring function (optional)
    """

    objective: OptimizationObjective = OptimizationObjective.MAXIMIZE_SHARPE
    max_iterations: int = 100
    n_jobs: int = -1
    early_stopping: bool = True
    early_stopping_patience: int = 10
    early_stopping_min_delta: float = 0.001
    validation_split: float = 0.2
    random_state: int | None = None
    custom_scorer: Callable[[BacktestResult], float] | None = None

    def __post_init__(self):
        """Validate configuration."""
        if self.max_iterations < 1:
            raise ValueError("max_iterations must be >= 1")
        if not 0 <= self.validation_split < 1:
            raise ValueError("validation_split must be in [0, 1)")
        if self.early_stopping_patience < 1:
            raise ValueError("early_stopping_patience must be >= 1")


@dataclass
class OptimizationResult:
    """Result of optimization.

    Attributes:
        best_params: Best parameters found
        best_score: Best score achieved
        best_result: BacktestResult with best parameters
        all_results: List of all (params, score, result) tuples
        optimization_time: Total optimization time in seconds
        n_iterations: Number of iterations performed
        converged: Whether optimization converged
        parameter_importance: Parameter importance scores (if available)
    """

    best_params: dict[str, Any]
    best_score: float
    best_result: BacktestResult | None
    all_results: list[tuple[dict[str, Any], float, BacktestResult | None]] = field(
        default_factory=list
    )
    optimization_time: float = 0.0
    n_iterations: int = 0
    converged: bool = False
    parameter_importance: dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert result to dictionary."""
        return {
            "best_params": self.best_params,
            "best_score": self.best_score,
            "best_result": self.best_result.to_dict() if self.best_result else None,
            "optimization_time": self.optimization_time,
            "n_iterations": self.n_iterations,
            "converged": self.converged,
            "parameter_importance": self.parameter_importance,
        }


class StrategyOptimizer(ABC):
    """Abstract base class for strategy optimizers.

    This class defines the interface for all optimization algorithms.

    Args:
        parameter_space: Parameter space to optimize over
        config: Optimizer configuration
    """

    def __init__(
        self,
        parameter_space: ParameterSpace,
        config: OptimizerConfig | None = None,
    ):
        """Initialize the optimizer."""
        self.parameter_space = parameter_space
        self.config = config or OptimizerConfig()
        self.logger = get_logger(f"optimizer.{self.__class__.__name__}")

        # Set random seed if provided
        if self.config.random_state is not None:
            np.random.seed(self.config.random_state)

        # Results storage
        self._results: list[tuple[dict[str, Any], float, BacktestResult | None]] = []
        self._best_score: float = float("-inf")
        self._best_params: dict[str, Any] | None = None
        self._best_result: BacktestResult | None = None

    def _score_result(self, result: BacktestResult) -> float:
        """Calculate score from backtest result.

        Args:
            result: Backtest result

        Returns:
            Score (higher is better)
        """
        if self.config.objective == OptimizationObjective.CUSTOM:
            if self.config.custom_scorer is None:
                raise ValueError("custom_scorer must be provided for CUSTOM objective")
            return self.config.custom_scorer(result)

        scores = {
            OptimizationObjective.MAXIMIZE_RETURN: result.total_return,
            OptimizationObjective.MAXIMIZE_SHARPE: result.sharpe_ratio,
            OptimizationObjective.MAXIMIZE_CALMAR: result.calmar_ratio,
            OptimizationObjective.MINIMIZE_DRAWDOWN: -result.max_drawdown,
            OptimizationObjective.MAXIMIZE_WIN_RATE: result.win_rate,
        }

        return scores.get(self.config.objective, result.sharpe_ratio)

    def _evaluate_params(
        self,
        params: dict[str, Any],
        backtest_fn: Callable[[dict[str, Any]], BacktestResult],
    ) -> tuple[float, BacktestResult]:
        """Evaluate parameters using backtest function.

        Args:
            params: Parameters to evaluate
            backtest_fn: Function that runs backtest with given parameters

        Returns:
            Tuple of (score, backtest_result)
        """
        try:
            result = backtest_fn(params)
            score = self._score_result(result)
            return score, result
        except Exception as e:
            self.logger.error(f"Error evaluating params {params}: {e}")
            return float("-inf"), None

    def _evaluate_params_parallel(
        self,
        param_list: list[dict[str, Any]],
        backtest_fn: Callable[[dict[str, Any]], BacktestResult],
        n_jobs: int = -1,
    ) -> list[tuple[dict[str, Any], float, BacktestResult]]:
        """Evaluate multiple parameters in parallel.

        Args:
            param_list: List of parameters to evaluate
            backtest_fn: Function that runs backtest with given parameters
            n_jobs: Number of parallel jobs

        Returns:
            List of (params, score, result) tuples
        """
        if n_jobs == -1:
            import os

            n_jobs = os.cpu_count() or 1

        results = []

        if n_jobs == 1 or len(param_list) == 1:
            # Sequential execution
            for params in param_list:
                score, result = self._evaluate_params(params, backtest_fn)
                results.append((params, score, result))
        else:
            # Parallel execution - note that backtest_fn must be serializable
            with ProcessPoolExecutor(max_workers=n_jobs) as executor:
                futures = {
                    executor.submit(self._evaluate_params, params, backtest_fn): params
                    for params in param_list
                }

                for future in as_completed(futures):
                    params = futures[future]
                    try:
                        score, result = future.result()
                        results.append((params, score, result))
                    except Exception as e:
                        self.logger.error(f"Error in parallel evaluation: {e}")
                        results.append((params, float("-inf"), None))

        return results

    def _update_best(
        self,
        params: dict[str, Any],
        score: float,
        result: BacktestResult | None,
    ) -> bool:
        """Update best parameters if score is better.

        Args:
            params: Parameters evaluated
            score: Score achieved
            result: Backtest result

        Returns:
            True if best was updated
        """
        if score > self._best_score:
            self._best_score = score
            self._best_params = params.copy()
            self._best_result = result
            return True
        return False

    def _check_early_stopping(self, iteration: int, no_improvement_count: int) -> bool:
        """Check if early stopping criteria are met.

        Args:
            iteration: Current iteration number
            no_improvement_count: Number of iterations without improvement

        Returns:
            True if should stop early
        """
        if not self.config.early_stopping:
            return False

        if no_improvement_count >= self.config.early_stopping_patience:
            self.logger.info(
                f"Early stopping at iteration {iteration}: "
                f"no improvement for {no_improvement_count} iterations"
            )
            return True

        return False

    @abstractmethod
    def optimize(
        self,
        backtest_fn: Callable[[dict[str, Any]], BacktestResult],
        callback: Callable[[int, dict[str, Any], float], None] | None = None,
    ) -> OptimizationResult:
        """Run optimization.

        Args:
            backtest_fn: Function that takes parameters and returns BacktestResult
            callback: Optional callback function (iteration, params, score)

        Returns:
            OptimizationResult with best parameters and all results
        """
        pass

    def get_best_params(self) -> dict[str, Any] | None:
        """Get best parameters found so far."""
        return self._best_params

    def get_best_score(self) -> float:
        """Get best score achieved so far."""
        return self._best_score

    def get_best_result(self) -> BacktestResult | None:
        """Get best backtest result."""
        return self._best_result
