"""GPyTorch-based Gaussian Process models with physics-informed priors."""

from typing import Callable, Optional, Tuple

import torch
from torch import Tensor
import gpytorch
from gpytorch.models import ExactGP
from gpytorch.means import ConstantMean, ZeroMean
from gpytorch.kernels import ScaleKernel, RBFKernel, MaternKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.distributions import MultivariateNormal
from gpytorch.mlls import ExactMarginalLogLikelihood

from botorch.models.gpytorch import GPyTorchModel
from botorch.posteriors.gpytorch import GPyTorchPosterior
from botorch.models.transforms.input import Normalize
from botorch.models.transforms.outcome import Standardize

from physics_informed_bo.models.base import SurrogateModel
from physics_informed_bo.models.physics_model import PhysicsMeanFunction


class _ExactGPModel(ExactGP, GPyTorchModel):
    """Core GPyTorch ExactGP model with BoTorch compatibility."""

    _num_outputs = 1

    def __init__(
        self,
        train_X: Tensor,
        train_y: Tensor,
        likelihood: GaussianLikelihood,
        mean_module: Optional[gpytorch.means.Mean] = None,
        kernel: str = "matern",
        ard_num_dims: Optional[int] = None,
    ):
        super().__init__(train_X, train_y.squeeze(-1), likelihood)

        self.mean_module = mean_module or ConstantMean()

        if kernel == "rbf":
            base_kernel = RBFKernel(ard_num_dims=ard_num_dims)
        elif kernel == "matern":
            base_kernel = MaternKernel(nu=2.5, ard_num_dims=ard_num_dims)
        else:
            raise ValueError(f"Unknown kernel: {kernel}. Use 'rbf' or 'matern'.")

        self.covar_module = ScaleKernel(base_kernel)

    def forward(self, X: Tensor) -> MultivariateNormal:
        mean = self.mean_module(X)
        covar = self.covar_module(X)
        return MultivariateNormal(mean, covar)


class StandardGP(SurrogateModel):
    """Standard Gaussian Process model (no physics, pure data-driven).

    Uses GPyTorch for the GP and is BoTorch-compatible for optimization.
    """

    def __init__(
        self,
        kernel: str = "matern",
        noise_variance: float = 0.01,
        learn_noise: bool = True,
        normalize_inputs: bool = True,
        standardize_outputs: bool = True,
        device: str = "cpu",
        dtype: torch.dtype = torch.float64,
    ):
        self.kernel = kernel
        self.noise_variance = noise_variance
        self.learn_noise = learn_noise
        self.normalize_inputs = normalize_inputs
        self.standardize_outputs = standardize_outputs
        self.device = torch.device(device)
        self.dtype = dtype
        self._model = None
        self._likelihood = None

    def fit(
        self,
        X: Tensor,
        y: Tensor,
        training_iterations: int = 100,
        lr: float = 0.1,
    ) -> None:
        """Fit the GP model by optimizing the marginal log likelihood."""
        X = X.to(device=self.device, dtype=self.dtype)
        y = y.to(device=self.device, dtype=self.dtype)

        if y.dim() == 1:
            y = y.unsqueeze(-1)

        self._likelihood = GaussianLikelihood()
        if not self.learn_noise:
            self._likelihood.noise = self.noise_variance
            self._likelihood.noise_covar.raw_noise.requires_grad_(False)

        self._model = _ExactGPModel(
            train_X=X,
            train_y=y,
            likelihood=self._likelihood,
            kernel=self.kernel,
            ard_num_dims=X.shape[-1],
        ).to(device=self.device, dtype=self.dtype)

        self._optimize_hyperparameters(X, y, training_iterations, lr)

    def _optimize_hyperparameters(
        self, X: Tensor, y: Tensor, n_iter: int, lr: float
    ) -> None:
        """Optimize GP hyperparameters via type-II MLE."""
        self._model.train()
        self._likelihood.train()

        optimizer = torch.optim.Adam(self._model.parameters(), lr=lr)
        mll = ExactMarginalLogLikelihood(self._likelihood, self._model)

        for _ in range(n_iter):
            optimizer.zero_grad()
            output = self._model(X)
            loss = -mll(output, y.squeeze(-1))
            loss.backward()
            optimizer.step()

        self._model.eval()
        self._likelihood.eval()

    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
        X = X.to(device=self.device, dtype=self.dtype)
        self._model.eval()
        self._likelihood.eval()

        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            posterior = self._likelihood(self._model(X))
            mean = posterior.mean.unsqueeze(-1)
            variance = posterior.variance.unsqueeze(-1)

        return mean, variance

    def posterior(self, X: Tensor):
        self._model.eval()
        self._likelihood.eval()
        return self._model.posterior(X)

    @property
    def model(self):
        """Access the underlying BoTorch-compatible GP model."""
        return self._model


class PhysicsInformedGP(SurrogateModel):
    """GP with a physics model as the mean function.

    The GP prior mean is set to the physics model predictions, so the GP
    learns the residual (discrepancy) between the physics model and reality.
    This is the core model of the platform.

    Architecture:
        f(x) = physics(x) + GP_residual(x)
        where GP_residual ~ GP(0, k(x, x'))
    """

    def __init__(
        self,
        physics_fn: Callable[[Tensor], Tensor],
        kernel: str = "matern",
        physics_output_scale: float = 1.0,
        learnable_physics_scale: bool = True,
        noise_variance: float = 0.01,
        learn_noise: bool = True,
        device: str = "cpu",
        dtype: torch.dtype = torch.float64,
    ):
        self.physics_fn = physics_fn
        self.kernel = kernel
        self.physics_output_scale = physics_output_scale
        self.learnable_physics_scale = learnable_physics_scale
        self.noise_variance = noise_variance
        self.learn_noise = learn_noise
        self.device = torch.device(device)
        self.dtype = dtype
        self._model = None
        self._likelihood = None

    def fit(
        self,
        X: Tensor,
        y: Tensor,
        training_iterations: int = 200,
        lr: float = 0.05,
    ) -> None:
        """Fit the physics-informed GP model."""
        X = X.to(device=self.device, dtype=self.dtype)
        y = y.to(device=self.device, dtype=self.dtype)

        if y.dim() == 1:
            y = y.unsqueeze(-1)

        self._likelihood = GaussianLikelihood()
        if not self.learn_noise:
            self._likelihood.noise = self.noise_variance
            self._likelihood.noise_covar.raw_noise.requires_grad_(False)

        physics_mean = PhysicsMeanFunction(
            physics_fn=self.physics_fn,
            output_scale=self.physics_output_scale,
            learnable_scale=self.learnable_physics_scale,
        )

        self._model = _ExactGPModel(
            train_X=X,
            train_y=y,
            likelihood=self._likelihood,
            mean_module=physics_mean,
            kernel=self.kernel,
            ard_num_dims=X.shape[-1],
        ).to(device=self.device, dtype=self.dtype)

        self._optimize_hyperparameters(X, y, training_iterations, lr)

    def _optimize_hyperparameters(
        self, X: Tensor, y: Tensor, n_iter: int, lr: float
    ) -> None:
        self._model.train()
        self._likelihood.train()

        optimizer = torch.optim.Adam(self._model.parameters(), lr=lr)
        mll = ExactMarginalLogLikelihood(self._likelihood, self._model)

        for _ in range(n_iter):
            optimizer.zero_grad()
            output = self._model(X)
            loss = -mll(output, y.squeeze(-1))
            loss.backward()
            optimizer.step()

        self._model.eval()
        self._likelihood.eval()

    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
        X = X.to(device=self.device, dtype=self.dtype)
        self._model.eval()
        self._likelihood.eval()

        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            posterior = self._likelihood(self._model(X))
            mean = posterior.mean.unsqueeze(-1)
            variance = posterior.variance.unsqueeze(-1)

        return mean, variance

    def posterior(self, X: Tensor):
        self._model.eval()
        self._likelihood.eval()
        return self._model.posterior(X)

    @property
    def model(self):
        return self._model

    def get_residuals(self, X: Tensor, y: Tensor) -> Tensor:
        """Compute residuals between physics predictions and observations."""
        with torch.no_grad():
            physics_pred = self.physics_fn(X)
        return y.squeeze() - physics_pred