"""Bayesian optimization for OpenClaw Trading.

This module provides Bayesian optimization using Gaussian Process modeling
with various acquisition functions.
"""

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,
    ParameterType,
    StrategyOptimizer,
)


class GaussianProcess:
    """Simple Gaussian Process implementation for Bayesian optimization.

    This is a lightweight implementation using the squared exponential kernel.
    For production use, consider using scikit-learn's GaussianProcessRegressor
    or GPy.

    Args:
        length_scale: Kernel length scale (default: 1.0)
        noise_level: Observation noise level (default: 1e-5)
    """

    def __init__(self, length_scale: float = 1.0, noise_level: float = 1e-5):
        """Initialize Gaussian Process."""
        self.length_scale = length_scale
        self.noise_level = noise_level
        self.X_train: np.ndarray | None = None
        self.y_train: np.ndarray | None = None
        self.K_inv: np.ndarray | None = None

    def _kernel(self, x1: np.ndarray, x2: np.ndarray) -> np.ndarray:
        """Squared exponential kernel."""
        sqdist = (
            np.sum(x1**2, axis=1).reshape(-1, 1)
            + np.sum(x2**2, axis=1)
            - 2 * np.dot(x1, x2.T)
        )
        return np.exp(-0.5 * sqdist / self.length_scale**2)

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """Fit the GP to training data.

        Args:
            X: Training inputs (n_samples, n_features)
            y: Training targets (n_samples,)
        """
        self.X_train = X.copy()
        self.y_train = y.copy()

        # Compute kernel matrix
        K = self._kernel(X, X) + self.noise_level * np.eye(len(X))

        # Compute inverse for prediction
        try:
            self.K_inv = np.linalg.inv(K)
        except np.linalg.LinAlgError:
            # Add more regularization if matrix is singular
            K += 1e-3 * np.eye(len(X))
            self.K_inv = np.linalg.inv(K)

    def predict(
        self, X: np.ndarray, return_std: bool = True
    ) -> tuple[np.ndarray, np.ndarray | None]:
        """Make predictions.

        Args:
            X: Test inputs (n_samples, n_features)
            return_std: Whether to return standard deviation

        Returns:
            Tuple of (mean, std) predictions
        """
        if self.X_train is None or self.K_inv is None:
            raise ValueError("GP not fitted yet")

        # Compute kernel between test and train
        K_s = self._kernel(X, self.X_train)

        # Predictive mean
        mu = K_s @ self.K_inv @ self.y_train

        if not return_std:
            return mu, None

        # Predictive variance
        K_ss = self._kernel(X, X)
        var = np.diag(K_ss) - np.sum(K_s @ self.K_inv * K_s, axis=1)
        var = np.maximum(var, 1e-10)  # Ensure positive
        std = np.sqrt(var)

        return mu, std


class AcquisitionFunction:
    """Acquisition functions for Bayesian optimization."""

    @staticmethod
    def ucb(
        mu: np.ndarray,
        std: np.ndarray,
        kappa: float = 2.0,
        maximize: bool = True,
    ) -> np.ndarray:
        """Upper Confidence Bound acquisition function.

        Args:
            mu: Predicted mean
            std: Predicted standard deviation
            kappa: Exploration parameter (higher = more exploration)
            maximize: Whether to maximize or minimize

        Returns:
            Acquisition values
        """
        if maximize:
            return mu + kappa * std
        return -mu + kappa * std

    @staticmethod
    def ei(
        mu: np.ndarray,
        std: np.ndarray,
        y_best: float,
        maximize: bool = True,
    ) -> np.ndarray:
        """Expected Improvement acquisition function.

        Args:
            mu: Predicted mean
            std: Predicted standard deviation
            y_best: Best observed value so far
            maximize: Whether to maximize or minimize

        Returns:
            Acquisition values
        """
        from scipy import stats

        if maximize:
            imp = mu - y_best
        else:
            imp = y_best - mu

        z = imp / (std + 1e-10)
        ei = imp * stats.norm.cdf(z) + std * stats.norm.pdf(z)
        return ei

    @staticmethod
    def pi(
        mu: np.ndarray,
        std: np.ndarray,
        y_best: float,
        xi: float = 0.01,
        maximize: bool = True,
    ) -> np.ndarray:
        """Probability of Improvement acquisition function.

        Args:
            mu: Predicted mean
            std: Predicted standard deviation
            y_best: Best observed value so far
            xi: Exploration parameter
            maximize: Whether to maximize or minimize

        Returns:
            Acquisition values
        """
        from scipy import stats

        if maximize:
            z = (mu - y_best - xi) / (std + 1e-10)
        else:
            z = (y_best - mu - xi) / (std + 1e-10)

        return stats.norm.cdf(z)


class BayesianOptimizer(StrategyOptimizer):
    """Bayesian optimizer using Gaussian Process.

    This optimizer uses Bayesian optimization with Gaussian Process modeling
    and acquisition functions to efficiently search the parameter space.

    Args:
        parameter_space: Parameter space to optimize over
        config: Optimizer configuration
        n_initial_points: Number of random initial points (default: 10)
        acquisition: Acquisition function ("ucb", "ei", "pi") (default: "ei")
        kappa: Exploration parameter for UCB (default: 2.0)
        xi: Exploration parameter for PI (default: 0.01)
        length_scale: GP kernel length scale (default: 1.0)
    """

    def __init__(
        self,
        parameter_space: ParameterSpace,
        config: OptimizerConfig | None = None,
        n_initial_points: int = 10,
        acquisition: str = "ei",
        kappa: float = 2.0,
        xi: float = 0.01,
        length_scale: float = 1.0,
    ):
        """Initialize Bayesian optimizer."""
        super().__init__(parameter_space, config)
        self.n_initial_points = n_initial_points
        self.acquisition = acquisition.lower()
        self.kappa = kappa
        self.xi = xi
        self.length_scale = length_scale

        # Validate acquisition function
        if self.acquisition not in ("ucb", "ei", "pi"):
            raise ValueError(f"Unknown acquisition function: {acquisition}")

        # GP model
        self.gp = GaussianProcess(length_scale=length_scale)

        # Parameter names for encoding/decoding
        self._param_names = list(parameter_space._parameters.keys())

    def _encode_params(self, params: dict[str, Any]) -> np.ndarray:
        """Encode parameters to numeric array."""
        encoded = []
        for name in self._param_names:
            param_range = self.parameter_space.get_parameter(name)
            value = params[name]

            if param_range.param_type == ParameterType.CATEGORICAL:
                # One-hot encoding
                choices = list(param_range.bounds)
                encoded.extend([1.0 if c == value else 0.0 for c in choices])
            else:
                encoded.append(float(value))

        return np.array(encoded)

    def _decode_params(self, encoded: np.ndarray) -> dict[str, Any]:
        """Decode numeric array to parameters."""
        params = {}
        idx = 0

        for name in self._param_names:
            param_range = self.parameter_space.get_parameter(name)

            if param_range.param_type == ParameterType.CATEGORICAL:
                # Decode one-hot
                choices = list(param_range.bounds)
                n_choices = len(choices)
                one_hot = encoded[idx : idx + n_choices]
                choice_idx = np.argmax(one_hot)
                params[name] = choices[choice_idx]
                idx += n_choices
            elif param_range.param_type == ParameterType.INTEGER:
                params[name] = int(round(encoded[idx]))
                idx += 1
            else:
                params[name] = encoded[idx]
                idx += 1

        return params

    def _acquisition_function(
        self, X: np.ndarray, y_best: float
    ) -> np.ndarray:
        """Evaluate acquisition function."""
        mu, std = self.gp.predict(X, return_std=True)

        if self.acquisition == "ucb":
            return AcquisitionFunction.ucb(mu, std, self.kappa)
        elif self.acquisition == "ei":
            return AcquisitionFunction.ei(mu, std, y_best)
        else:  # pi
            return AcquisitionFunction.pi(mu, std, y_best, self.xi)

    def _suggest_next_point(self) -> dict[str, Any]:
        """Suggest next point to evaluate using acquisition function."""
        # Sample random candidates
        n_candidates = 1000
        candidates = [
            self.parameter_space.sample_random() for _ in range(n_candidates)
        ]
        X_candidates = np.array([self._encode_params(c) for c in candidates])

        # Evaluate acquisition function
        acq_values = self._acquisition_function(X_candidates, self._best_score)

        # Select best candidate
        best_idx = np.argmax(acq_values)
        return candidates[best_idx]

    def optimize(
        self,
        backtest_fn: Callable[[dict[str, Any]], BacktestResult],
        callback: Callable[[int, dict[str, Any], float], None] | None = None,
    ) -> OptimizationResult:
        """Run Bayesian 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 Bayesian optimization")

        results = []
        no_improvement_count = 0
        previous_best = float("-inf")

        # Initial random sampling
        n_initial = min(self.n_initial_points, self.config.max_iterations)
        self.logger.info(f"Running {n_initial} initial random samples")

        for iteration in range(n_initial):
            params = self.parameter_space.sample_random()
            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)

            if score > previous_best:
                previous_best = score

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

        # Bayesian optimization loop
        self.logger.info("Starting Bayesian optimization loop")

        for iteration in range(n_initial, self.config.max_iterations):
            # Prepare training data
            X_train = np.array([self._encode_params(r[0]) for r in results])
            y_train = np.array([r[1] for r in results])

            # Normalize y for numerical stability
            y_mean = np.mean(y_train)
            y_std = np.std(y_train) + 1e-10
            y_train_normalized = (y_train - y_mean) / y_std

            # Fit GP model
            try:
                self.gp.fit(X_train, y_train_normalized)
            except Exception as e:
                self.logger.warning(f"GP fitting failed: {e}. Using random sample.")
                params = self.parameter_space.sample_random()
            else:
                # Suggest next point
                params = self._suggest_next_point()

            # Evaluate suggested point
            score, result = self._evaluate_params(params, backtest_fn)

            results.append((params, score, result))
            improved = self._update_best(params, score, result)

            if callback:
                callback(iteration, params, score)

            # Early stopping check
            if improved:
                if self._best_score - previous_best < self.config.early_stopping_min_delta:
                    no_improvement_count += 1
                else:
                    no_improvement_count = 0
                previous_best = self._best_score
            else:
                no_improvement_count += 1

            if self._check_early_stopping(iteration, no_improvement_count):
                break

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

        optimization_time = time.time() - start_time
        converged = no_improvement_count < self.config.early_stopping_patience

        # Calculate parameter importance
        parameter_importance = self._calculate_parameter_importance(results)

        self.logger.info(
            f"Bayesian optimization 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=converged,
            parameter_importance=parameter_importance,
        )

    def _calculate_parameter_importance(
        self, results: list[tuple[dict[str, Any], float, BacktestResult]]
    ) -> dict[str, float]:
        """Calculate parameter importance based on variance in results."""
        if len(results) < 2:
            return {}

        importance = {}

        for name in self._param_names:
            param_range = self.parameter_space.get_parameter(name)
            values = [r[0][name] for r in results]
            scores = [r[1] for r in results]

            if param_range.param_type in (
                ParameterType.CONTINUOUS,
                ParameterType.INTEGER,
            ):
                # Calculate correlation between parameter value and score
                numeric_values = np.array([float(v) for v in values])
                if np.std(numeric_values) > 1e-10:
                    correlation = np.corrcoef(numeric_values, scores)[0, 1]
                    importance[name] = abs(correlation)
                else:
                    importance[name] = 0.0
            else:
                # For categorical/discrete, calculate variance of scores per value
                value_scores: dict[Any, list[float]] = {}
                for v, s in zip(values, scores, strict=False):
                    if v not in value_scores:
                        value_scores[v] = []
                    value_scores[v].append(s)

                # Higher variance between groups = more important
                group_means = [np.mean(scores) for scores in value_scores.values()]
                if len(group_means) > 1:
                    importance[name] = np.std(group_means)
                else:
                    importance[name] = 0.0

        # Normalize importance scores
        total = sum(importance.values())
        if total > 0:
            importance = {k: v / total for k, v in importance.items()}

        return importance
