import inspect
from typing import Any

import torch
from torch import Tensor, Size

from gpytorch.models import ApproximateGP
from gpytorch.variational import (
    CholeskyVariationalDistribution,
    VariationalStrategy,
    IndependentMultitaskVariationalStrategy,
)
from gpytorch.means import (
    ZeroMean,
    ConstantMean,
    Mean,
)
from gpytorch.kernels import (
    Kernel,
    ScaleKernel,
    RBFKernel,
    MaternKernel,
    RQKernel,
    RFFKernel,
    PeriodicKernel,
    LinearKernel,
)
from gpytorch.distributions import MultivariateNormal
from gpytorch.likelihoods import (
    Likelihood,
    GaussianLikelihood,
    BernoulliLikelihood,
    LaplaceLikelihood,
    SoftmaxLikelihood,
    StudentTLikelihood,
    BetaLikelihood,
    MultitaskGaussianLikelihood,
)


KERNEL_REGISTRY: dict[str, type[Kernel]] = {
    "rq": RQKernel,
    "matern": MaternKernel,
    "rbf": RBFKernel,
    "rff": RFFKernel,
    "periodic": PeriodicKernel,
    "linear": LinearKernel,
}
MEAN_REGISTRY: dict[str, type[Mean]] = {
    "constant": ConstantMean,
    "zero": ZeroMean,
}
LIKELIHOOD_REGISTRY: dict[str, type[Likelihood]] = {
    "gaussian": GaussianLikelihood,
    "bernoulli": BernoulliLikelihood,
    "laplace": LaplaceLikelihood,
    "softmax": SoftmaxLikelihood,
    "studentt": StudentTLikelihood,
    "beta": BetaLikelihood,
}


def _instantiate_from_spec(
    spec: str | type | Any,
    expected_base: type,
    registry: dict[str, type],
    kwargs: dict[str, Any],
    kind: str,
) -> Any:
    if isinstance(spec, str):
        cls = registry.get(spec.lower())
        if cls is None:
            raise ValueError(f"Unknown {kind}: {spec}")
    elif isinstance(spec, expected_base):
        return spec
    elif inspect.isclass(spec):
        if not issubclass(spec, expected_base):
            raise ValueError(f"{kind} class must inherit from {expected_base.__name__}")
        cls = spec
    else:
        raise ValueError(f"{kind} must be a string id, {expected_base.__name__} class, or {expected_base.__name__} instance")

    accepted_kwargs: dict[str, Any] = {}
    try:
        signature = inspect.signature(cls.__init__)
        has_var_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in signature.parameters.values())
        for key, value in kwargs.items():
            if has_var_kwargs or key in signature.parameters:
                accepted_kwargs[key] = value
    except (TypeError, ValueError):
        accepted_kwargs = kwargs

    try:
        return cls(**accepted_kwargs)
    except TypeError as exc:
        raise ValueError(f"Could not instantiate {kind} class '{cls.__name__}' with kwargs {accepted_kwargs}") from exc


def _get_kernel(kernel: str | type[Kernel] | Kernel, n_inputs: int, use_ard: bool, batch_shape: Size) -> Kernel:
    return _instantiate_from_spec(
        spec=kernel,
        expected_base=Kernel,
        registry=KERNEL_REGISTRY,
        kwargs={
            "ard_num_dims": n_inputs if use_ard else None,
            "batch_shape": batch_shape,
            "num_samples": n_inputs,
        },
        kind="kernel",
    )


def _get_mean_module(mean: str | type[Mean] | Mean, batch_shape: Size) -> Mean:
    return _instantiate_from_spec(
        spec=mean,
        expected_base=Mean,
        registry=MEAN_REGISTRY,
        kwargs={"batch_shape": batch_shape},
        kind="mean",
    )


def _get_likelihood(likelihood: str | type[Likelihood] | Likelihood, n_outputs: int) -> Likelihood:
    if n_outputs > 1:
        if isinstance(likelihood, str):
            return MultitaskGaussianLikelihood(num_tasks=n_outputs)

    return _instantiate_from_spec(
        spec=likelihood,
        expected_base=Likelihood,
        registry=LIKELIHOOD_REGISTRY,
        kwargs={"num_tasks": n_outputs},
        kind="likelihood",
    )


class GPRegressor(ApproximateGP):
    def __init__(self, *,
                 n_inducing_points: int,
                 n_inputs: int,
                 n_outputs: int,
                 kernel: str | type[Kernel] | Kernel = "rq",
                 mean: str | type[Mean] | Mean = "constant",
                 likelihood: str | type[Likelihood] | Likelihood = "gaussian",
                 use_ard: bool = True,
                 ) -> None:
        '''
        Arguments
        ---------
        n_inducing_points: int
            Number of inducing points
        n_inputs: int
            Number of input dimensions
        n_outputs: int
            Number of output dimensions
        kernel: str | type[gpytorch.kernels.Kernel] | gpytorch.kernels.Kernel
            Kernel to use as string id, class, or instance.
        mean: str | type[gpytorch.means.Mean] | gpytorch.means.Mean
            Mean to use as string id, class, or instance.
        likelihood: str | type[gpytorch.likelihoods.Likelihood] | gpytorch.likelihoods.Likelihood
            Likelihood to use as string id, class, or instance.
        use_ard: bool
            Whether to use Automatic Relevance Determination
        '''
        batch_shape = Size([]) if n_outputs == 1 else Size([n_outputs])
        inducing_points_shape = (n_inducing_points, n_inputs) if n_outputs == 1 else (n_outputs, n_inducing_points, n_inputs)
        inducing_points: Tensor = torch.randn(inducing_points_shape, dtype=torch.float32)

        variational_distribution = CholeskyVariationalDistribution(
            n_inducing_points,
            batch_shape=batch_shape,
        )
        base_variational_strategy = VariationalStrategy(
            self,
            inducing_points,
            variational_distribution,
            learn_inducing_locations=True,
        )
        variational_strategy = base_variational_strategy if n_outputs == 1 else IndependentMultitaskVariationalStrategy(
            base_variational_strategy=base_variational_strategy,
            num_tasks=n_outputs,
        )

        super().__init__(variational_strategy)
        self.mean_module = _get_mean_module(mean, batch_shape)
        self.covar_module = ScaleKernel(_get_kernel(kernel, n_inputs, use_ard, batch_shape), batch_shape=batch_shape)
        self.likelihood = _get_likelihood(likelihood, n_outputs)

    def forward(self, x: Tensor) -> MultivariateNormal:
        mean_x: Tensor = self.mean_module(x)  # type: ignore
        covar_x: Tensor = self.covar_module(x)  # type: ignore

        latent_pred: MultivariateNormal = MultivariateNormal(mean_x, covar_x)

        return latent_pred
