"""Bayesian optimization for strategy parameter tuning.

Implements a simple Gaussian Process surrogate with Expected Improvement
acquisition -- all from scratch using only the Python standard library
(math module).  Designed for low-dimensional search spaces (2-8 params)
typical of prediction-market strategy tuning.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import Any, Callable

from horizon.backtest import backtest
from horizon.result import BacktestResult


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class BayesianOptResult:
    """Result of Bayesian optimization.

    Attributes:
        best_params: Parameter dict that achieved the highest objective score.
        best_score: Best objective value found.
        all_trials: List of dicts, each with ``params`` (dict) and ``score`` (float).
        n_trials: Total number of trials evaluated.
        convergence: Best score seen up to and including each trial index.
    """

    best_params: dict[str, float] = field(default_factory=dict)
    best_score: float = float("-inf")
    all_trials: list[dict[str, Any]] = field(default_factory=list)
    n_trials: int = 0
    convergence: list[float] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Gaussian Process (internal)
# ---------------------------------------------------------------------------

def _rbf_kernel(x1: list[float], x2: list[float], length_scale: float) -> float:
    """RBF (squared-exponential) kernel: k(x, x') = exp(-||x-x'||^2 / (2*l^2))."""
    sq_dist = sum((a - b) ** 2 for a, b in zip(x1, x2))
    return math.exp(-sq_dist / (2.0 * length_scale ** 2))


def _kernel_matrix(
    X: list[list[float]],
    length_scale: float,
    noise: float = 1e-6,
) -> list[list[float]]:
    """Compute the kernel matrix K + noise * I for a set of points *X*."""
    n = len(X)
    K = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(i, n):
            val = _rbf_kernel(X[i], X[j], length_scale)
            K[i][j] = val
            K[j][i] = val
        K[i][i] += noise
    return K


def _cholesky(A: list[list[float]]) -> list[list[float]]:
    """Cholesky decomposition A = L L^T.

    Adds progressively larger jitter on the diagonal when the matrix is
    near-singular to guarantee numerical stability.
    """
    n = len(A)
    for attempt in range(10):
        jitter = 10.0 ** (-8 + attempt) if attempt > 0 else 0.0
        L = [[0.0] * n for _ in range(n)]
        success = True
        for i in range(n):
            for j in range(i + 1):
                s = sum(L[i][k] * L[j][k] for k in range(j))
                if i == j:
                    diag = A[i][i] + jitter - s
                    if diag <= 0.0:
                        success = False
                        break
                    L[i][j] = math.sqrt(diag)
                else:
                    if L[j][j] == 0.0:
                        success = False
                        break
                    L[i][j] = (A[i][j] - s) / L[j][j]
            if not success:
                break
        if success:
            return L
    # Fallback: return diagonal matrix (effectively ignoring correlations)
    L = [[0.0] * n for _ in range(n)]
    for i in range(n):
        L[i][i] = math.sqrt(abs(A[i][i]) + 1e-4)
    return L


def _solve_triangular_lower(L: list[list[float]], b: list[float]) -> list[float]:
    """Solve L x = b where L is lower-triangular (forward substitution)."""
    n = len(b)
    x = [0.0] * n
    for i in range(n):
        s = sum(L[i][j] * x[j] for j in range(i))
        if L[i][i] == 0.0:
            x[i] = 0.0
        else:
            x[i] = (b[i] - s) / L[i][i]
    return x


def _solve_triangular_upper(U: list[list[float]], b: list[float]) -> list[float]:
    """Solve U x = b where U is upper-triangular (back substitution)."""
    n = len(b)
    x = [0.0] * n
    for i in range(n - 1, -1, -1):
        s = sum(U[j][i] * x[j] for j in range(i + 1, n))
        if U[i][i] == 0.0:
            x[i] = 0.0
        else:
            x[i] = (b[i] - s) / U[i][i]
    return x


def _cholesky_solve(L: list[list[float]], b: list[float]) -> list[float]:
    """Solve (L L^T) x = b via forward then back substitution."""
    y = _solve_triangular_lower(L, b)
    # L^T is the transpose of L -- stored as the same nested list but indexed [j][i]
    return _solve_triangular_upper(L, y)


class _GP:
    """Minimal Gaussian Process regressor with RBF kernel.

    Designed for the small datasets (10-100 points) typical of Bayesian
    optimization, so O(N^3) Cholesky is perfectly fine.
    """

    def __init__(self, length_scale: float = 1.0, noise: float = 1e-6) -> None:
        self.length_scale = length_scale
        self.noise = noise
        self._X: list[list[float]] = []
        self._y: list[float] = []
        self._L: list[list[float]] = []
        self._alpha: list[float] = []
        self._y_mean: float = 0.0

    def fit(self, X: list[list[float]], y: list[float]) -> None:
        """Fit GP on observed data."""
        self._X = X
        self._y_mean = sum(y) / len(y) if y else 0.0
        self._y = [yi - self._y_mean for yi in y]
        K = _kernel_matrix(X, self.length_scale, self.noise)
        self._L = _cholesky(K)
        self._alpha = _cholesky_solve(self._L, self._y)

    def predict(self, x: list[float]) -> tuple[float, float]:
        """Predict mean and standard deviation at a single point *x*.

        Returns:
            (mean, std) -- both finite floats (std >= 0).
        """
        if not self._X:
            return self._y_mean, 1.0

        k_star = [_rbf_kernel(x, xi, self.length_scale) for xi in self._X]

        # Posterior mean
        mu = sum(a * k for a, k in zip(self._alpha, k_star)) + self._y_mean

        # Posterior variance
        v = _solve_triangular_lower(self._L, k_star)
        k_ss = _rbf_kernel(x, x, self.length_scale)
        var = k_ss - sum(vi * vi for vi in v)
        var = max(var, 0.0)
        std = math.sqrt(var)

        # Ensure finite
        if not math.isfinite(mu):
            mu = self._y_mean
        if not math.isfinite(std):
            std = 0.0

        return mu, std


# ---------------------------------------------------------------------------
# Acquisition function
# ---------------------------------------------------------------------------

def _normal_cdf(x: float) -> float:
    """Standard normal CDF using the error function."""
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))


def _normal_pdf(x: float) -> float:
    """Standard normal PDF."""
    return math.exp(-0.5 * x * x) / math.sqrt(2.0 * math.pi)


def _expected_improvement(mu: float, std: float, f_best: float, xi: float = 0.01) -> float:
    """Expected Improvement acquisition function.

    EI(x) = (mu - f_best - xi) * Phi(z) + std * phi(z)
    where z = (mu - f_best - xi) / std.

    Returns 0.0 when std is effectively zero.
    """
    if std < 1e-12:
        return 0.0
    z = (mu - f_best - xi) / std
    ei = (mu - f_best - xi) * _normal_cdf(z) + std * _normal_pdf(z)
    return max(ei, 0.0)


# ---------------------------------------------------------------------------
# Helper: estimate GP length scale from the search space
# ---------------------------------------------------------------------------

def _estimate_length_scale(search_space: dict[str, tuple[float, float]]) -> float:
    """Heuristic: use the average range of the search space dimensions."""
    if not search_space:
        return 1.0
    ranges = [hi - lo for lo, hi in search_space.values()]
    avg_range = sum(ranges) / len(ranges) if ranges else 1.0
    # A length scale of ~1/3 the average range works well empirically
    return max(avg_range / 3.0, 1e-6)


# ---------------------------------------------------------------------------
# Helper: normalize / denormalize parameters to [0, 1]
# ---------------------------------------------------------------------------

def _normalize_params(
    params: dict[str, float],
    search_space: dict[str, tuple[float, float]],
) -> list[float]:
    """Map a parameter dict to a [0,1]^d vector (order follows search_space keys)."""
    vec = []
    for name, (lo, hi) in search_space.items():
        val = params.get(name, lo)
        if hi - lo > 0:
            vec.append((val - lo) / (hi - lo))
        else:
            vec.append(0.5)
    return vec


def _denormalize_vec(
    vec: list[float],
    search_space: dict[str, tuple[float, float]],
) -> dict[str, float]:
    """Map a [0,1]^d vector back to a parameter dict."""
    params: dict[str, float] = {}
    for i, (name, (lo, hi)) in enumerate(search_space.items()):
        params[name] = lo + vec[i] * (hi - lo)
    return params


# ---------------------------------------------------------------------------
# Acquisition maximizer (random + local refinement)
# ---------------------------------------------------------------------------

def _maximize_acquisition(
    gp: _GP,
    f_best: float,
    search_space: dict[str, tuple[float, float]],
    n_candidates: int = 500,
    rng: random.Random | None = None,
    xi: float = 0.01,
) -> dict[str, float]:
    """Find the point that maximizes Expected Improvement.

    Uses a two-stage approach:
    1. Evaluate EI at *n_candidates* random points.
    2. Do a simple perturbation-based local search around the best candidate.
    """
    if rng is None:
        rng = random.Random()

    dim = len(search_space)
    best_ei = -1.0
    best_vec: list[float] = [0.5] * dim

    # Stage 1: random candidates (in [0,1]^d normalized space)
    for _ in range(n_candidates):
        vec = [rng.random() for _ in range(dim)]
        mu, std = gp.predict(vec)
        ei = _expected_improvement(mu, std, f_best, xi=xi)
        if ei > best_ei:
            best_ei = ei
            best_vec = vec

    # Stage 2: local refinement (simple random perturbation)
    step_size = 0.1
    for _ in range(50):
        candidate = [
            max(0.0, min(1.0, v + rng.gauss(0, step_size)))
            for v in best_vec
        ]
        mu, std = gp.predict(candidate)
        ei = _expected_improvement(mu, std, f_best, xi=xi)
        if ei > best_ei:
            best_ei = ei
            best_vec = candidate
        step_size *= 0.95  # anneal

    return _denormalize_vec(best_vec, search_space)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def bayesian_optimize(
    data: Any,
    pipeline_factory: Callable[..., list[Callable]],
    search_space: dict[str, tuple[float, float]],
    objective: str = "sharpe_ratio",
    n_trials: int = 50,
    n_initial: int = 10,
    markets: list[str] | None = None,
    risk: Any = None,
    seed: int | None = None,
    **backtest_kwargs: Any,
) -> BayesianOptResult:
    """Run Bayesian optimization over strategy parameters.

    The optimizer first evaluates *n_initial* random parameter samples, then
    uses a Gaussian Process surrogate model with Expected Improvement to
    choose the most promising points for the remaining trials.

    Args:
        data: Historical data (same formats accepted by ``hz.backtest``).
        pipeline_factory: ``factory(params) -> pipeline list``.  Called once
            per trial with the candidate parameter dict.
        search_space: Mapping of parameter names to ``(min, max)`` bounds.
        objective: Attribute name on :class:`~horizon.analytics.Metrics` to
            maximize (e.g. ``"sharpe_ratio"``, ``"total_return_pct"``).
        n_trials: Total number of parameter evaluations.
        n_initial: Number of initial random samples before GP-guided search.
        markets: Passed through to ``backtest()``.
        risk: Passed through to ``backtest()``.
        seed: Random seed for reproducibility.
        **backtest_kwargs: Extra keyword arguments forwarded to ``backtest()``.

    Returns:
        :class:`BayesianOptResult` with the best parameters, scores, and
        convergence trace.
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    if n_trials < 1:
        return BayesianOptResult(n_trials=0)

    n_initial = max(1, min(n_initial, n_trials))

    rng = random.Random(seed)

    all_trials: list[dict[str, Any]] = []
    convergence: list[float] = []
    best_score = float("-inf")
    best_params: dict[str, float] = {}

    # Collect normalized X and raw y for the GP
    X_observed: list[list[float]] = []
    y_observed: list[float] = []

    # Estimate a reasonable GP length scale from the search space
    ls = _estimate_length_scale(search_space)
    # Work in normalized [0,1] space, so rescale length scale
    ls_norm = 1.0 / 3.0  # ~1/3 of [0,1] range

    for trial_idx in range(n_trials):
        if trial_idx < n_initial:
            # Random sample
            params = {
                name: lo + rng.random() * (hi - lo)
                for name, (lo, hi) in search_space.items()
            }
        else:
            # GP-guided sample
            gp = _GP(length_scale=ls_norm, noise=1e-4)
            gp.fit(X_observed, y_observed)
            params = _maximize_acquisition(
                gp,
                f_best=best_score,
                search_space=search_space,
                n_candidates=max(200, 100 * len(search_space)),
                rng=rng,
            )

        # Evaluate candidate via backtest
        try:
            pipeline = pipeline_factory(params)
            bt_result: BacktestResult = backtest(
                data=data,
                pipeline=pipeline,
                markets=markets,
                risk=risk,
                **backtest_kwargs,
            )
            score = getattr(bt_result.metrics, objective, 0.0)
            if score is None or not math.isfinite(score):
                score = 0.0
        except Exception:
            score = 0.0

        # Record
        all_trials.append({"params": params, "score": score})
        X_observed.append(_normalize_params(params, search_space))
        y_observed.append(score)

        if score > best_score:
            best_score = score
            best_params = dict(params)

        convergence.append(best_score)

    return BayesianOptResult(
        best_params=best_params,
        best_score=best_score,
        all_trials=all_trials,
        n_trials=len(all_trials),
        convergence=convergence,
    )
