"""Grid search optimizer for OpenClaw Trading.

This module provides grid search optimization for strategy parameters.
"""

import time
from collections.abc import Callable
from typing import Any

import numpy as np

from openclaw.backtest.engine import BacktestResult
from openclaw.optimizer.base import (
    OptimizationResult,
    OptimizerConfig,
    ParameterSpace,
    StrategyOptimizer,
)


class GridSearchOptimizer(StrategyOptimizer):
    """Grid search optimizer.

    This optimizer exhaustively searches through a manually specified subset
    of the hyperparameter space. It evaluates all combinations of parameters
    in the grid.

    Args:
        parameter_space: Parameter space to optimize over
        config: Optimizer configuration
        n_points: Number of points per continuous dimension (default: 5)
    """

    def __init__(
        self,
        parameter_space: ParameterSpace,
        config: OptimizerConfig | None = None,
        n_points: int = 5,
    ):
        """Initialize grid search optimizer."""
        super().__init__(parameter_space, config)
        self.n_points = n_points

    def optimize(
        self,
        backtest_fn: Callable[[dict[str, Any]], BacktestResult],
        callback: Callable[[int, dict[str, Any], float], None] | None = None,
    ) -> OptimizationResult:
        """Run grid search 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
        """
        start_time = time.time()
        self.logger.info("Starting grid search optimization")

        # Generate all parameter combinations
        param_grid = self.parameter_space.get_grid_points(self.n_points)
        total_combinations = len(param_grid)

        self.logger.info(f"Total parameter combinations: {total_combinations}")

        if total_combinations == 0:
            raise ValueError("No parameter combinations to evaluate")

        if total_combinations > self.config.max_iterations:
            self.logger.warning(
                f"Parameter grid ({total_combinations}) exceeds max_iterations "
                f"({self.config.max_iterations}). Truncating grid."
            )
            # Sample from the grid to respect max_iterations
            indices = np.linspace(0, total_combinations - 1, self.config.max_iterations, dtype=int)
            param_grid = [param_grid[i] for i in indices]
            total_combinations = len(param_grid)

        # Evaluate all combinations
        n_jobs = self.config.n_jobs
        results = []

        # Check if we can use parallel execution
        if n_jobs == 1 or total_combinations == 1:
            # Sequential execution
            for iteration, params in enumerate(param_grid):
                score, result = self._evaluate_params(params, backtest_fn)
                results.append((params, score, result))
                self._update_best(params, score, result)

                if callback:
                    callback(iteration, params, score)

                self.logger.debug(
                    f"Iteration {iteration + 1}/{total_combinations}: "
                    f"params={params}, score={score:.4f}"
                )
        else:
            # Parallel execution
            results = self._evaluate_params_parallel(param_grid, backtest_fn, n_jobs)

            # Update best and call callback
            for iteration, (params, score, result) in enumerate(results):
                self._update_best(params, score, result)

                if callback:
                    callback(iteration, params, score)

                self.logger.debug(
                    f"Iteration {iteration + 1}/{total_combinations}: "
                    f"params={params}, score={score:.4f}"
                )

        optimization_time = time.time() - start_time

        self.logger.info(
            f"Grid search complete. Best score: {self._best_score:.4f}, "
            f"Best params: {self._best_params}"
        )

        return OptimizationResult(
            best_params=self._best_params or {},
            best_score=self._best_score,
            best_result=self._best_result,
            all_results=results,
            optimization_time=optimization_time,
            n_iterations=len(results),
            converged=True,  # Grid search always "converges" by checking all points
        )

    def get_grid_size(self) -> int:
        """Get the total number of parameter combinations in the grid.

        Returns:
            Number of combinations
        """
        return len(self.parameter_space.get_grid_points(self.n_points))
