import jax
import jax.numpy as jnp
import jax.random
from jax import jit
from flax.linen.module import Module, compact
from flax.linen import Dense
import flax.linen
from dataclasses import field
from collections.abc import Callable
from typing import (Any, Callable, Iterable, Optional, Tuple, Union)

PRNGKey = Any
Shape = Iterable[int]
Dtype = Any
Array = Any

def linear(x,layer_params=None,unit_params=None):
  return x

def relu(x,layer_params=None,unit_params=None):
  return flax.linen.relu(x)

def arcsinh(x,layer_params=None,unit_params=None):
  return jnp.arcsinh(x)

def tanh(x,layer_params=None,unit_params=None):
  return jax.nn.tanh(x)

def sigma(x,layer_params=None,unit_params=None):
  return flax.linen.sigmoid(x)

def prelu(x,layer_params=None,unit_params=None):
  return relu(x) - unit_params*relu(-x)

def zero_mean_func(x, mean_params):
  sample_num = x.shape[0]
  return jnp.zeros((sample_num,1))

def constant_mean_func(x, mean_params):
  sample_num = x.shape[0]
  constant = mean_params
  return constant*jnp.ones((sample_num,1))

def square_exp_kernel_func(x1, x2, kernel_params):
  std = jnp.exp(kernel_params[0])
  length_scale = jnp.exp(kernel_params[1])
  diff = jnp.expand_dims(x1/length_scale, axis=1) - \
         jnp.expand_dims(x2/length_scale, axis=0)
  kernel_mat = std**2*jnp.exp(-jnp.sum(diff**2/2, axis=2))
  return kernel_mat

def rbf_kernel_func(x1, x2, kernel_params):
  std = jnp.exp(kernel_params[0])
  length_scale = jnp.expand_dims(jnp.exp(kernel_params[1:]), 0)
  diff = jnp.expand_dims(x1/length_scale, axis=1) - \
         jnp.expand_dims(x2/length_scale, axis=0)
  kernel_mat = std**2*jnp.exp(-jnp.sum(diff**2/2, axis=2))
  return kernel_mat

def normal_initialize(key, shape, dtype=jnp.float32, mean=0.0, std=1.0):
  return std * jax.random.normal(key, shape, dtype) + mean

class Activation(Module):
  """
  """
  activation_func: Callable[[Array, Array, Array], Array] = linear
  activation_per_layer_params_num: int = 0
  activation_per_unit_params_num: int = 0
  unit_num: int = 0
  layer_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  unit_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  dtype: Any = jnp.float32
  layer_params: Array = field(init=False)
  unit_params: Array = field(init=False)


  def setup(self):
    self.layer_params = self.param('per_layer', self.layer_init, (self.activation_per_layer_params_num,))
    self.unit_params = self.param('per_unit', self.unit_init, (self.activation_per_unit_params_num,self.unit_num))
  
  def __call__(self, x: Array) -> Array:
    """
    """
    x = jnp.asarray(x, self.dtype)
    y = self.activation_func(x,self.layer_params,self.unit_params)
    return y

class Mean(Module):
  """
  """
  mean_func: Callable[[Array, Array], Array] = zero_mean_func
  mean_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  mean_params_num: int = 0
  dtype: Any = jnp.float32
  mean_params: Array = field(init=False)

  def setup(self):
    """
    """
    self.mean_params = self.param('mean', self.mean_init, (self.mean_params_num,))
  
  def __call__(self, x: Array) -> Array:
    """
    """
    x = jnp.asarray(x, self.dtype)
    m = self.mean_func(x, self.mean_params)
    return m

class Kernel(Module):
  """
  """
  kernel_func: Callable[[Array, Array, Array], Array]  = square_exp_kernel_func
  kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  kernel_params_num: int = 2
  dtype: Any = jnp.float32
  kernel_params: Array = field(init=False)

  def setup(self):
    """
    """
    self.kernel_params = self.param('kernel', self.kernel_init, (self.kernel_params_num,))
  
  def __call__(self, x1: Array, x2: Array) -> Array:
    """
    """
    x1 = jnp.asarray(x1, self.dtype)
    x2 = jnp.asarray(x2, self.dtype)
    K = self.kernel_func(x1, x2, self.kernel_params)
    return K

class GP(Module):
  """
  """
  mean_func: Callable[[Array, Array], Array]  = zero_mean_func
  kernel_func: Callable[[Array, Array, Array], Array]  = square_exp_kernel_func
  mean_params_num: int = 0
  kernel_params_num: int = 2
  mean_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  likelihood_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  dtype: Any = jnp.float32
  mean: Any = field(init=False)
  kernel: Any = field(init=False)
  likelihood_std: Array = field(init=False)
  

  def setup(self):
    """
    """
    submodule = Mean(mean_func = self.mean_func, 
                     mean_params_num = self.mean_params_num, 
                     mean_init = self.mean_init)
    self.mean = submodule
    submodule = Kernel(kernel_func = self.kernel_func, 
                         kernel_params_num = self.kernel_params_num, 
                         kernel_init = self.kernel_init)
    self.kernel = submodule
    self.likelihood_std = self.param('likelihood_std', 
                                      self.likelihood_init, 
                                      (1,))
  
  def __call__(self, x_data: Array, y_data: Array) -> Array:
    """
    """
    return self.post_mean(x_pred=x_data, x_data=x_data, y_data=y_data)

  def sample(self, random_key: PRNGKey, x_pred: Array, x_data: Array, y_data: Array, n_sample: int) -> Array:
    """
    """
    x_pred = jnp.asarray(x_pred, self.dtype)
    x_data = jnp.asarray(x_data, self.dtype)
    y_data = jnp.asarray(y_data, self.dtype)
    pred_dims = (x_pred.shape[0], n_sample)
    gaussian_mat = jax.random.normal(random_key, shape = pred_dims)
    cov_mat = self.post_cov(x_pred, x_data, y_data)
    mean_vec = self.post_mean(x_pred)
    lower_chol_cov_mat = jnp.linalg.cholesky(cov_mat).T
    samples = mean_vec + lower_chol_cov_mat @ gaussian_mat
    return samples
  
  def likelihood(self, x_data: Array) -> Array:
    """
    """
    batch_size = x_data.shape[0]
    likelihood_var = jnp.exp(self.likelihood_std)**2
    eye_mat = jnp.eye(batch_size)
    likelihood_mat = likelihood_var * eye_mat
    return likelihood_mat

  def prior_mean(self, x_data: Array) -> Array:
    """
    """
    x_data = jnp.asarray(x_data, self.dtype)
    m = self.mean(x_data)
    return m

  def prior_cov(self, x_data_1: Array, x_data_2: Array) -> Array:
    """
    """
    x_data_1 = jnp.asarray(x_data_1, self.dtype)
    x_data_2 = jnp.asarray(x_data_2, self.dtype)
    K = self.kernel(x_data_1,x_data_2)
    return K
    
  def post_mean(self, x_pred: Array, x_data: Array, y_data: Array) -> Array:
    """
    """
    x_pred = jnp.asarray(x_pred, self.dtype)
    x_data = jnp.asarray(x_data, self.dtype)
    y_data = jnp.asarray(y_data, self.dtype)
    mean_data_vec = self.prior_mean(x_data)
    mean_pred_vec = self.prior_mean(x_pred)
    cov_pred_data_mat = self.prior_cov(x_pred, x_data)
    cov_data_data_mat = self.prior_cov(x_data, x_data)
    likelihood_mat = self.likelihood(x_data)
    diff_vec = y_data - mean_data_vec
    K_mat = cov_data_data_mat + likelihood_mat
    Kinvdiff = jnp.linalg.solve(K_mat,diff_vec)
    mean_post_pred_vec = mean_pred_vec + cov_pred_data_mat@Kinvdiff
    return mean_post_pred_vec

  def post_cov(self, x_pred: Array, x_data: Array) -> Array:
    """
    """
    x_pred = jnp.asarray(x_pred, self.dtype)
    x_data = jnp.asarray(x_data, self.dtype)
    cov_data_pred_mat = self.prior_cov(x_data, x_pred)
    cov_data_data_mat = self.prior_cov(x_data, x_data)
    cov_pred_pred_mat = self.prior_cov(x_pred, x_pred)
    likelihood_mat = self.likelihood(x_data)
    K_mat = cov_data_data_mat + likelihood_mat
    Kinvk = jnp.linalg.solve(K_mat,cov_data_pred_mat)
    cov_post_pred_pred_mat = cov_pred_pred_mat - cov_data_pred_mat.T@Kinvk
    return cov_post_pred_pred_mat

  def post_std(self, x_pred: Array, x_data: Array) -> Array:
    """
    """
    x_pred = jnp.asarray(x_pred, self.dtype)
    x_data = jnp.asarray(x_data, self.dtype)
    cov_data_pred_mat = self.prior_cov(x_data, x_pred)
    cov_data_data_mat = self.prior_cov(x_data, x_data)
    likelihood_mat = self.likelihood(x_data)
    K_mat = cov_data_data_mat + likelihood_mat
    Kinvk = jnp.linalg.solve(K_mat,cov_data_pred_mat)
    prior_var_pred = lambda x_pred: self.prior_cov(jnp.expand_dims(x_pred,0), jnp.expand_dims(x_pred,0))
    var_pred = lambda x_pred, l_mat, r_mat: (prior_var_pred(x_pred) - jnp.dot(l_mat.T, r_mat))[0,0]
    var_pred_vmap = jax.vmap(var_pred, (0, 1, 1), 0)
    var_post_pred_vec = jnp.expand_dims(var_pred_vmap(x_pred, cov_data_pred_mat, Kinvk), 1)
    std_post_pred_vec = jnp.sqrt(var_post_pred_vec)
    return std_post_pred_vec

  def negative_loglikelihood_loss(self, x_data: Array, y_data: Array):
    """
    """
    x_data = jnp.asarray(x_data, self.dtype)
    y_data = jnp.asarray(y_data, self.dtype)
    mean_data_vec = self.prior_mean(x_data)
    cov_data_data_mat = self.prior_cov(x_data,x_data)
    likelihood_mat = self.likelihood(x_data)
    K = likelihood_mat + cov_data_data_mat
    diff_vec = (y_data-mean_data_vec)
    yKinvy = diff_vec.T @ (jnp.linalg.solve(K,diff_vec))
    (_, ld) = jnp.linalg.slogdet(K)
    loss = jnp.abs(yKinvy[0,0]) + ld
    return loss

  def minibatch_negative_loglikelihood_loss(self, key: PRNGKey, batch_size: int, x_data: Array, y_data: Array):
    """
    """
    x_data = jnp.asarray(x_data, self.dtype)
    y_data = jnp.asarray(y_data, self.dtype)
    arangemat = jnp.arange(0,x_data.shape[0])
    sample = jax.random.choice(key,arangemat,(batch_size,),replace=False)
    x_batch = x_data[sample,:]
    y_batch = y_data[sample,:]
    loss = self.negative_loglikelihood_loss(x_batch, y_batch)
    return loss

  def get_params(self):
    return self.params

class NNDRGP(GP):
  """
  """
  mean_func: Callable[[Array, Array], Array]  = zero_mean_func
  kernel_func: Callable[[Array, Array, Array], Array]  = square_exp_kernel_func
  activation_func: Callable = linear
  mean_params_num: int = 0
  kernel_params_num: int = 2
  activation_per_layer_params_num: int = 0
  activation_per_unit_params_num: int = 0
  layer_num: int = 1
  output_num: int = 1
  hidden_num: int = 1
  mean_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  likelihood_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  layer_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  unit_init: Callable[[PRNGKey, Shape, Dtype], Array] = normal_initialize
  dtype: Any = jnp.float32
  mean: Any = field(init=False)
  kernel: Any = field(init=False)
  dr: Any = field(init=False)
  likelihood_std: Array = field(init=False)

  def setup(self):
    """
    """
    super().setup()
    submodule = {}
    for i in range(self.layer_num):
      layername = "layer"+str(i)
      submodule[layername] = Dense(self.hidden_num)
    layername = "layer"+str(self.layer_num)
    submodule[layername] = Dense(self.output_num)
    self.dr = submodule
    submodule = {}
    for i in range(self.layer_num):
      layername = "layer"+str(i)
      submodule[layername] = Activation(activation_func=self.activation_func,
                                        activation_per_layer_params_num=self.activation_per_layer_params_num,
                                        activation_per_unit_params_num=self.activation_per_unit_params_num,
                                        unit_num=self.hidden_num,
                                        layer_init=self.layer_init,
                                        unit_init=self.unit_init)
    self.activation = submodule
  
  def __call__(self, x_data: Array, y_data: Array) -> Array:
    """
    """
    return self.post_mean(x_pred=x_data, x_data=x_data, y_data=y_data)
  
  def dim_reduce(self, x: Array):
    """
    """
    x = jnp.asarray(x, self.dtype)
    y = x
    for i in range(self.layer_num):
      layername = 'layer'+str(i)
      y = self.dr[layername](y)
      y = self.activation[layername](y)
    layername = 'layer'+str(self.layer_num)
    y = self.dr[layername](y)
    return y

  def sample(self, random_key: PRNGKey, x_pred: Array, x_data: Array, y_data: Array, n_sample: int) -> Array:
    """
    """
    return super().sample(random_key, x_pred, x_data, y_data, n_sample)
  
  def likelihood(self, x_data: Array) -> Array:
    """
    """
    x_data_reduced = self.dim_reduce(x_data)
    return super().likelihood(x_data_reduced)

  def prior_mean(self, x_data: Array) -> Array:
    """
    """
    x_data_reduced = self.dim_reduce(x_data)
    return super().prior_mean(x_data_reduced)

  def prior_cov(self, x_data_1: Array, x_data_2: Array) -> Array:
    """
    """
    x_data_1_reduced = self.dim_reduce(x_data_1)
    x_data_2_reduced = self.dim_reduce(x_data_2)
    return super().prior_cov(x_data_1_reduced,x_data_2_reduced)
    
  def post_mean(self, x_pred: Array, x_data: Array, y_data: Array) -> Array:
    """
    """
    return super().post_mean(x_pred,x_data,y_data)

  def post_cov(self, x_pred: Array, x_data: Array) -> Array:
    """
    """
    return super().post_cov(x_pred,x_data)

  def post_std(self, x_pred: Array, x_data: Array) -> Array:
    """
    """
    return super().post_std(x_pred,x_data)

  def negative_loglikelihood_loss(self, x_data: Array, y_data: Array):
    return super().negative_loglikelihood_loss(x_data,y_data)

  def minibatch_negative_loglikelihood_loss(self, key: PRNGKey, batch_size: int, x_data: Array, y_data: Array):
    return super().minibatch_negative_loglikelihood_loss(key,batch_size,x_data,y_data)

  def get_params(self):
    return super().get_params()