"""Hyperparameter optimization for Gaussian Process models.

Provides functionality to optimize GP hyperparameters (sigma2, ell, period) by
maximizing the log marginal likelihood of the training data.

The log marginal likelihood is:
    log p(y|X,θ) = -0.5 * y^T K^{-1} y - 0.5 * log|K| - n/2 * log(2π)

where K is the kernel matrix (including noise) and θ are the hyperparameters.
"""

import numpy as np
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Callable
from scipy import optimize

from .types import Observations


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

    Attributes:
        sigma2: Optimized kernel variance (must be > 0).
        ell: Optimized length scale (must be > 0).
        period: Optimized period (for periodic kernels only, None otherwise).
        log_likelihood: Log marginal likelihood at optimum.
        success: Whether optimization succeeded.
        scipy_result: Raw scipy optimization result object.
    """

    sigma2: float
    ell: float
    period: Optional[float]
    log_likelihood: float
    success: bool
    scipy_result: Any


# =============================================================================
# Log Marginal Likelihood Computation
# =============================================================================

# Import C implementation
from ._core import log_marginal_likelihood as _log_marginal_likelihood_c


def compute_log_marginal_likelihood(
    obs: Observations,
    sigma2: float,
    ell: float,
    period: float,
    kernel: str,
    ell_p: Optional[float] = None,
) -> float:
    """Compute log marginal likelihood: log p(y|X,θ).

    Formula: -0.5 * y^T K^{-1} y - 0.5 * log|K| - n/2 * log(2π)

    This is a thin wrapper around the C implementation.

    Args:
        obs: Observations object with training data.
        sigma2: Kernel variance parameter.
        ell: SE decay length scale (used by RBF, Matérn, locally-periodic).
        period: Period parameter (for periodic kernels).
        kernel: Kernel type.
        ell_p: Periodic length scale (for periodic/locally-periodic kernels).
            Defaults to ell if not specified.

    Returns:
        Log marginal likelihood value (more positive is better).

    Raises:
        ValueError: If parameters are invalid or kernel doesn't support observations.
        RuntimeError: If C implementation is not available.
    """
    return _log_marginal_likelihood_c(
        obs, sigma2=sigma2, ell=ell, ell_p=ell_p, period=period, kernel=kernel
    )


# =============================================================================
# Optimization Functions
# =============================================================================


def _create_objective_function(
    obs: Observations, kernel: str
) -> Tuple[Callable, int]:
    """Create objective function for optimization.

    Args:
        obs: Observations object.
        kernel: Kernel type.

    Returns:
        Tuple of (objective_function, n_params) where objective_function
        takes a parameter vector and returns negative log likelihood.

    Raises:
        ValueError: If kernel is unsupported or incompatible with observations.
    """
    # Validate kernel early
    valid_kernels = ["rbf", "matern52", "matern32", "periodic", "locally_periodic"]
    if kernel not in valid_kernels:
        raise ValueError(
            f"Unknown kernel: '{kernel}'. "
            f"Supported: {', '.join(valid_kernels)}"
        )

    if obs.x_g is not None:
        raise ValueError(
            "Integral observations (g) are not currently supported for optimization. "
            "Please use only f and/or h observations."
        )

    # Determine number of parameters
    needs_period = kernel in ["periodic", "locally_periodic"]
    n_params = 3 if needs_period else 2

    def objective(params: np.ndarray) -> float:
        """Negative log likelihood (for minimization)."""
        sigma2 = params[0]
        ell = params[1]
        period = params[2] if needs_period else 1.0

        # Safeguard: ensure positive parameters
        if sigma2 <= 0 or ell <= 0:
            return 1e10
        if needs_period and period <= 0:
            return 1e10

        try:
            log_lik = compute_log_marginal_likelihood(obs, sigma2, ell, period, kernel)
            return -log_lik  # Minimize negative log likelihood
        except (ValueError, RuntimeError):
            return 1e10  # Return large value on error

    return objective, n_params


def _get_default_bounds(kernel: str) -> Dict[str, Tuple[float, float]]:
    """Get default parameter bounds for optimization.

    Args:
        kernel: Kernel type.

    Returns:
        Dictionary mapping parameter names to (min, max) bounds.
    """
    bounds = {
        "sigma2": (1e-4, 100.0),
        "ell": (1e-4, 10.0),
    }

    if kernel in ["periodic", "locally_periodic"]:
        bounds["period"] = (1e-2, 100.0)

    return bounds


def fit(
    obs: Observations,
    kernel: str = "rbf",
    bounds: Optional[Dict[str, Tuple[float, float]]] = None,
    method: str = "differential_evolution",
    x0: Optional[Dict[str, float]] = None,
    **optimizer_kwargs,
) -> OptimizationResult:
    """Optimize GP hyperparameters by maximizing log marginal likelihood.

    This function finds the hyperparameters (sigma2, ell, and optionally period)
    that maximize the log marginal likelihood of the observed data. It supports
    both global optimization methods (recommended for non-convex likelihood
    surfaces) and local optimization methods (faster but require good initial guesses).

    Args:
        obs: Observations object with x_f, y_f (and optionally x_h, y_h).
            Note: Integral observations (g) are not currently supported.
        kernel: Kernel type. One of:
            - 'rbf': RBF/Squared Exponential kernel
            - 'matern52': Matérn 5/2 kernel (twice differentiable)
            - 'matern32': Matérn 3/2 kernel (once mean-square differentiable)
            - 'periodic': Periodic kernel (requires period parameter)
            - 'locally_periodic': Locally-periodic kernel (requires period parameter)
            Default: 'rbf'.
        bounds: Parameter bounds as a dict mapping parameter names to (min, max) tuples.
            Example: {'sigma2': (0.01, 10.0), 'ell': (0.1, 5.0)}.
            If None, uses sensible defaults. Default: None.
        method: Optimization method. Options:
            Global methods (recommended):
            - 'differential_evolution': Genetic algorithm (default, robust)
            - 'dual_annealing': Simulated annealing variant
            - 'shgo': Simplicial homology global optimization
            - 'basinhopping': Basin-hopping algorithm

            Local methods (require x0):
            - 'L-BFGS-B': Quasi-Newton method (fast, requires gradients)
            - 'Nelder-Mead': Simplex method (derivative-free)
            - 'Powell': Powell's method (derivative-free)
            - 'SLSQP': Sequential Least Squares Programming
            - Other scipy.optimize.minimize methods

            Default: 'differential_evolution'.
        x0: Initial parameter values (required for local methods, ignored for global methods).
            Example: {'sigma2': 1.0, 'ell': 1.0} or {'sigma2': 1.0, 'ell': 1.0, 'period': 5.0}.
            Default: None.
        **optimizer_kwargs: Additional keyword arguments passed to the optimizer.
            For differential_evolution: maxiter, popsize, atol, tol, seed, etc.
            For minimize: maxiter, maxfev, ftol, gtol, etc.
            See scipy.optimize documentation for details.

    Returns:
        OptimizationResult with fields:
            - sigma2: Optimized kernel variance
            - ell: Optimized length scale
            - period: Optimized period (None for non-periodic kernels)
            - log_likelihood: Log marginal likelihood at optimum
            - success: Whether optimization succeeded
            - scipy_result: Raw scipy optimization result

    Raises:
        ValueError: If kernel is unsupported, bounds are invalid, or observations
            contain unsupported types (e.g., g for any kernel, h for matern32).
        TypeError: If obs is not an Observations instance.

    Example:
        Basic usage with default optimizer:

        >>> import gp4c
        >>> import numpy as np
        >>> x = np.linspace(0, 10, 50)
        >>> y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        >>> obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)
        >>> result = gp4c.fit(obs, kernel='rbf')
        >>> print(f"Optimal: sigma2={result.sigma2:.3f}, ell={result.ell:.3f}")

        Using L-BFGS-B with custom bounds:

        >>> result = gp4c.fit(
        ...     obs,
        ...     kernel='rbf',
        ...     method='L-BFGS-B',
        ...     x0={'sigma2': 1.0, 'ell': 1.0},
        ...     bounds={'sigma2': (0.1, 5.0), 'ell': (0.1, 3.0)}
        ... )

        Optimizing periodic kernel:

        >>> result = gp4c.fit(
        ...     obs,
        ...     kernel='periodic',
        ...     bounds={'period': (0.5, 2.0)},
        ...     method='differential_evolution',
        ...     maxiter=200,
        ...     seed=42
        ... )

        Using derivative observations:

        >>> x_train = np.linspace(0, 5, 20)
        >>> y_train = np.sin(x_train)
        >>> dy_train = np.cos(x_train)
        >>> obs = gp4c.Observations(
        ...     x_f=x_train, y_f=y_train, noise_f=0.01,
        ...     x_h=x_train, y_h=dy_train, noise_h=0.01
        ... )
        >>> result = gp4c.fit(obs, kernel='matern52')

    Note:
        - The likelihood surface can be non-convex with multiple local optima.
          Global methods are recommended for initial exploration.
        - Derivative observations (h) can significantly improve optimization,
          especially for Matérn 5/2 and RBF kernels.
        - For periodic data, you can constrain the period search range via bounds.
        - Optimization can be sensitive to the choice of bounds and initial values.
    """
    # Validate inputs
    if not isinstance(obs, Observations):
        raise TypeError("obs must be an Observations instance")

    # Get default bounds and merge with user-provided bounds
    default_bounds = _get_default_bounds(kernel)
    if bounds is not None:
        default_bounds.update(bounds)
    bounds = default_bounds

    # Create objective function
    objective, n_params = _create_objective_function(obs, kernel)

    # Convert bounds dict to list of tuples in correct order
    param_names = ["sigma2", "ell"]
    if kernel in ["periodic", "locally_periodic"]:
        param_names.append("period")

    bounds_list = [bounds[name] for name in param_names]

    # Choose optimization method
    global_methods = ["differential_evolution", "dual_annealing", "shgo", "basinhopping"]

    if method in global_methods:
        # Global optimization
        if method == "differential_evolution":
            result = optimize.differential_evolution(
                objective, bounds_list, **optimizer_kwargs
            )
        elif method == "dual_annealing":
            result = optimize.dual_annealing(objective, bounds_list, **optimizer_kwargs)
        elif method == "shgo":
            result = optimize.shgo(objective, bounds_list, **optimizer_kwargs)
        elif method == "basinhopping":
            # Basin-hopping needs an initial point
            if x0 is None:
                # Use midpoint of bounds
                x0_array = np.array([(b[0] + b[1]) / 2.0 for b in bounds_list])
            else:
                x0_array = np.array([x0[name] for name in param_names])

            minimizer_kwargs = {"method": "L-BFGS-B", "bounds": bounds_list}
            result = optimize.basinhopping(
                objective, x0_array, minimizer_kwargs=minimizer_kwargs, **optimizer_kwargs
            )
        else:
            raise ValueError(f"Unknown global optimization method: {method}")

    else:
        # Local optimization with scipy.optimize.minimize
        if x0 is None:
            raise ValueError(f"Local optimization method '{method}' requires x0 (initial guess)")

        # Convert x0 dict to array
        x0_array = np.array([x0[name] for name in param_names])

        # Validate x0 is within bounds
        for i, (name, val) in enumerate(zip(param_names, x0_array)):
            lb, ub = bounds_list[i]
            if not (lb <= val <= ub):
                raise ValueError(
                    f"Initial value for {name} ({val}) is outside bounds ({lb}, {ub})"
                )

        result = optimize.minimize(
            objective, x0_array, method=method, bounds=bounds_list, **optimizer_kwargs
        )

    # Extract optimized parameters
    opt_params = result.x
    sigma2_opt = opt_params[0]
    ell_opt = opt_params[1]
    period_opt = opt_params[2] if n_params == 3 else None

    # Compute log likelihood at optimum (remember objective is negative log likelihood)
    log_likelihood_opt = -result.fun

    # Create result object
    return OptimizationResult(
        sigma2=sigma2_opt,
        ell=ell_opt,
        period=period_opt,
        log_likelihood=log_likelihood_opt,
        success=result.success,
        scipy_result=result,
    )
