from jax.config import config

config.update("jax_enable_x64", True)

from dataclasses import dataclass
from typing import Dict

from jax import jit
import jax.numpy as jnp
import jax.random as jr
import jax.scipy as jscipy

# import types
from beartype.typing import Union
from jaxtyping import Float
from gpjax.typing import (
    Array,
    ScalarFloat,
)
from simple_pytree import static_field
import tensorflow_probability.substrates.jax as tfp
import gpjax as gpx
from gpjax.base.param import param_field


@dataclass
class SharedKernel(gpx.kernels.AbstractKernel):
    r"""Kernel that is shared between first half and second half of data, but no correlation between first and second half of data."""
    lengthscale: Union[ScalarFloat, Float[Array, " D"]] = param_field(
        jnp.array(1.0), bijector=tfb.Softplus()
    )
    variance: ScalarFloat = param_field(
            jnp.array(1.0), bijector=tfb.Softplus()
    )
    tmax: float = 1.0
    name: str = "Shared"
    def __call__(self, x1: Float[Array, " D"], x2: Float[Array, " D"]) -> ScalarFloat:
        l = self.lengthscale
        v = self.variance
        tmax = self.tmax

        if x1 <= xstar and x2 > xstar:
            K = 0 # no correlation between first half of data and second
        else:
            K = v * jnp.exp(-((x1 - x2) / l)**2)
        return K.squeeze()