from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import numpy as np
    from gpytorch.likelihoods import GaussianLikelihood

    from .turbo_gp import TurboGP


def standardize_y(y) -> tuple[float, float]:
    import numpy as np

    y_array = np.asarray(y, dtype=float)
    center = float(np.median(y_array))
    scale = float(np.std(y_array))
    if not np.isfinite(scale) or scale <= 0.0:
        scale = 1.0
    return center, scale


def fit_gp(
    x_obs_list: list,
    y_obs_list: list,
    num_dim: int,
    *,
    num_steps: int = 50,
) -> tuple[
    "TurboGP | None",
    "GaussianLikelihood | None",
    float,
    float,
]:
    import numpy as np
    import torch
    from gpytorch.constraints import Interval
    from gpytorch.likelihoods import GaussianLikelihood
    from gpytorch.mlls import ExactMarginalLogLikelihood

    from .turbo_gp import TurboGP

    x = np.asarray(x_obs_list, dtype=float)
    y = np.asarray(y_obs_list, dtype=float)
    n = x.shape[0]
    if n == 0:
        return None, None, 0.0, 1.0
    if n == 1:
        gp_y_mean = float(y[0])
        gp_y_std = 1.0
        return None, None, gp_y_mean, gp_y_std
    gp_y_mean, gp_y_std = standardize_y(y)
    y_centered = y - gp_y_mean
    z = y_centered / gp_y_std
    train_x = torch.as_tensor(x, dtype=torch.float64)
    train_y = torch.as_tensor(z, dtype=torch.float64)
    noise_constraint = Interval(5e-4, 0.2)
    lengthscale_constraint = Interval(0.005, 2.0)
    outputscale_constraint = Interval(0.05, 20.0)
    likelihood = GaussianLikelihood(noise_constraint=noise_constraint).to(
        dtype=train_y.dtype
    )
    model = TurboGP(
        train_x=train_x,
        train_y=train_y,
        likelihood=likelihood,
        lengthscale_constraint=lengthscale_constraint,
        outputscale_constraint=outputscale_constraint,
        ard_dims=num_dim,
    ).to(dtype=train_x.dtype)
    model.covar_module.outputscale = torch.tensor(1.0, dtype=train_x.dtype)
    model.covar_module.base_kernel.lengthscale = torch.full(
        (num_dim,), 0.5, dtype=train_x.dtype
    )
    likelihood.noise = torch.tensor(0.005, dtype=train_y.dtype)
    model.train()
    likelihood.train()
    mll = ExactMarginalLogLikelihood(likelihood, model)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    for _ in range(num_steps):
        optimizer.zero_grad()
        output = model(train_x)
        loss = -mll(output, train_y)
        loss.backward()
        optimizer.step()
    model.eval()
    likelihood.eval()
    return model, likelihood, gp_y_mean, gp_y_std


def latin_hypercube(num_points: int, num_dim: int, *, rng) -> np.ndarray:
    import numpy as np

    x = np.zeros((num_points, num_dim))
    centers = (1.0 + 2.0 * np.arange(0.0, num_points)) / float(2 * num_points)
    for j in range(num_dim):
        x[:, j] = centers[rng.permutation(num_points)]
    pert = rng.uniform(-1.0, 1.0, size=(num_points, num_dim)) / float(2 * num_points)
    x += pert
    return x


def argmax_random_tie(values, *, rng) -> int:
    import numpy as np

    if values.ndim != 1:
        raise ValueError(values.shape)
    max_val = float(np.max(values))
    idx = np.nonzero(values >= max_val)[0]
    if idx.size == 0:
        return int(rng.integers(values.size))
    if idx.size == 1:
        return int(idx[0])
    j = int(rng.integers(idx.size))
    return int(idx[j])


def pareto_front(mu, se) -> np.ndarray:
    import numpy as np

    if mu.shape != se.shape or mu.ndim != 1:
        raise ValueError((mu.shape, se.shape))
    n = mu.size
    if n == 0:
        return np.zeros((0,), dtype=bool)
    order = np.argsort(-mu)
    mu_sorted = mu[order]
    se_sorted = se[order]
    is_pareto_sorted = np.zeros_like(mu_sorted, dtype=bool)
    best_se = float("inf")
    for i in range(n):
        if se_sorted[i] < best_se:
            best_se = float(se_sorted[i])
            is_pareto_sorted[i] = True
    is_pareto = np.zeros_like(is_pareto_sorted, dtype=bool)
    is_pareto[order] = is_pareto_sorted
    return is_pareto


def arms_from_pareto_fronts(x_cand, mu, se, num_arms, rng) -> np.ndarray:
    import numpy as np

    if x_cand.ndim != 2:
        raise ValueError(x_cand.shape)
    if mu.shape != se.shape or mu.ndim != 1:
        raise ValueError((mu.shape, se.shape))
    if mu.size != x_cand.shape[0]:
        raise ValueError((mu.size, x_cand.shape[0]))

    mu_neg = -mu

    i = np.argsort(mu_neg)
    x_cand_sorted = x_cand[i]
    se_sorted = se[i]

    i_all = list(range(len(se_sorted)))
    i_keep = []

    while len(i_keep) < num_arms:
        se_max = float("-inf")
        i_front = []
        for idx in i_all:
            if se_sorted[idx] >= se_max:
                i_front.append(idx)
                se_max = se_sorted[idx]
        if len(i_keep) + len(i_front) <= num_arms:
            i_keep.extend(i_front)
        else:
            i_keep.extend(
                rng.choice(i_front, size=num_arms - len(i_keep), replace=False)
            )
        i_all = sorted(set(i_all) - set(i_front))

    i_keep = np.array(i_keep)
    x_arms = x_cand_sorted[i_keep]

    return x_arms


def sobol_perturb_np(
    x_center, lb, ub, num_candidates, mask, *, sobol_engine
) -> np.ndarray:
    import numpy as np

    sobol_samples = sobol_engine.random(num_candidates)
    lb_array = np.asarray(lb)
    ub_array = np.asarray(ub)
    pert = lb_array + (ub_array - lb_array) * sobol_samples
    candidates = np.tile(x_center, (num_candidates, 1))
    if np.any(mask):
        candidates[mask] = pert[mask]
    return candidates


def raasp(
    x_center, lb, ub, num_candidates, *, num_pert: int = 20, rng, sobol_engine
) -> np.ndarray:
    import numpy as np

    num_dim = x_center.shape[-1]
    prob_perturb = min(num_pert / num_dim, 1.0)
    mask = rng.random((num_candidates, num_dim)) <= prob_perturb
    ind = np.where(np.sum(mask, axis=1) == 0)[0]
    if len(ind) > 0:
        mask[ind, rng.integers(0, num_dim, size=len(ind))] = True
    return sobol_perturb_np(
        x_center, lb, ub, num_candidates, mask, sobol_engine=sobol_engine
    )


def to_unit(x, bounds) -> np.ndarray:
    import numpy as np

    lb = bounds[:, 0]
    ub = bounds[:, 1]
    if np.any(ub <= lb):
        raise ValueError(bounds)
    return (x - lb) / (ub - lb)


def from_unit(x_unit, bounds) -> np.ndarray:
    lb = bounds[:, 0]
    ub = bounds[:, 1]
    return lb + x_unit * (ub - lb)
