import tensorflow as tf

from rltf.tf_utils import tf_inv


class BLR(tf.layers.Layer):
  """Bayesian Linear Regression Layer
  By default:
    - Uses tf.float64 for internal variables since higher precision is required for inverting matrices
    - Inverts the precision matrix using the woodburry formula - more stable and efficient
    -
  NOTE:
    - Internal variables are not captured by `tf.trainable_variables()` and by optimizers
      - This is because `self.trainable_variables` returns an empty list
      - The fact that variables are set as trainable at instantiation does not affect this
    - Internal variables are captured when selecting variables from a given scope
      - This means that if an agent and a target network with identical structures contain BLR layers,
        the BLR variables will get updated, if all variables within the relevant scopes are captured
    - For training variables, one needs to use the op returned by `self.train()`
  """

  def __init__(self, tau, sigma_e, mode="mean", w_dim=None, bias=False, dtype=tf.float64, name=None):

    super().__init__(trainable=False, dtype=dtype, name=name)

    assert mode in ["mean", "ts"]

    self.sigma  = sigma_e
    self.beta   = 1.0 / self.sigma**2
    self.tau    = tau
    self.bias   = bias
    self.w_dim  = w_dim
    self.mode   = mode

    # Custom TF Tensors and Ops
    self.w_mu     = None
    self.w_Sigma  = None
    self.w_Lambda = None
    self.w        = None  # Sampled value for w when using Thompson Sampling
    self.reset_op = None  # Reset all trainable variables to their inNeed to make sure it is itial values

    self.input_spec = tf.layers.InputSpec(min_ndim=2, max_ndim=2)


  def build(self, input_shape):
    input_shape = tf.TensorShape(input_shape)

    if input_shape[-1].value is None:
      raise ValueError('The last dimension of the inputs to `BLR` should be defined. Found `None`.')

    if self.w_dim is None:
      self.w_dim      = input_shape[-1].value
      self.input_spec = tf.layers.InputSpec(min_ndim=2, max_ndim=2, axes={-1: self.w_dim})
    else:
      assert self.w_dim == input_shape[-1].value

    I = tf.eye(self.w_dim, dtype=self.dtype)

    mu_init     = tf.zeros([self.w_dim, 1], dtype=self.dtype)
    Sigma_init  = 1.0/self.tau * I
    Lambda_init = self.tau * I

    self.w_mu     = self.add_variable("w_mu",
                                      shape=[self.w_dim, 1],
                                      # initializer=lambda *args, **kwargs: mu_init,
                                      initializer=tf.zeros_initializer,
                                      trainable=True)

    self.w_Sigma  = self.add_variable("w_Sigma",
                                      shape=[self.w_dim, self.w_dim],
                                      initializer=lambda *args, **kwargs: Sigma_init,
                                      trainable=True)

    self.w_Lambda = self.add_variable("w_Lambda",
                                      shape=[self.w_dim, self.w_dim],
                                      initializer=lambda *args, **kwargs: Lambda_init,
                                      trainable=True)

    self.w        = self.add_variable("w",
                                      shape=[self.w_dim, 1],
                                      # initializer=lambda *args, **kwargs: mu_init,
                                      initializer=tf.zeros_initializer,
                                      trainable=False)

    # Build the reset op
    self.reset_op = self._tf_update_params(mu_init, Sigma_init, Lambda_init)

    # Add debug histogrrams
    # tf.summary.histogram("debug/BLR/w_mu",      self.w_mu)
    # tf.summary.histogram("debug/BLR/w_Sigma",   self.w_Sigma)
    # tf.summary.histogram("debug/BLR/w_Lambda",  self.w_Lambda)
    # tf.summary.histogram("debug/BLR/w_ts",      self.w)

    self.built = True


  def call(self, inputs, **kwargs):
    """ Compute the posterior predictive distribution
    Args:
      X: tf.Tensor, `shape=[None, D]`. The feature matrix
    Returns:
      List of `tf.Tensor`s:
        mu: tf.Tensor, `shape=[None, 1]. The mean at each test point
        var: tf.Tensor, `shape=[None, 1]. The variance at each test point
    """
    X = self._cast_input(inputs)

    # Thompson Sampling Output
    if self.mode == "ts":
      mu = tf.matmul(X, self.w)
    # Bayesian Regression Output
    else:
      mu = tf.matmul(X, self.w_mu)

    # var ends up being diag(sigma**2 + matmul(matmul(X, w_Sigma), X.T))
    var = self.sigma**2 + tf.reduce_sum(tf.matmul(X, self.w_Sigma) * X, axis=-1, keepdims=True)
    # std = tf.sqrt(var)

    outputs = [mu, var]
    outputs = [self._cast_output(t) for t in outputs]
    return outputs


  def train(self, X, y):
    """Compute the weight posteriror of Bayesian Linear Regression
    Args:
      X: tf.Tensor, `shape=[None, D]`. The feature matrix
      y: tf.Tensor, `shape=[None, 1]`. The correct outputs
    Returns:
      tf.Op which performs the update operation
    """
    X = self._cast_input(X)
    y = self._cast_input(y)

    # Compute the posterior precision matrix
    w_Lambda = self.w_Lambda + self.beta * tf.matmul(X, X, transpose_a=True)

    # Compute the posterior covariance matrix
    X_norm  = 1.0 / self.sigma * X
    w_Sigma = tf_inv.woodburry_inverse(self.w_Sigma, tf.transpose(X_norm), X_norm)

    error = tf.losses.mean_squared_error(tf.matmul(w_Lambda, w_Sigma), tf.eye(self.w_dim))
    tf.summary.scalar("debug/BLR/inv_error", error)

    # Compute the posterior mean
    w_mu = tf.matmul(w_Sigma, self.beta * tf.matmul(X, y, True) + tf.matmul(self.w_Lambda, self.w_mu))

    return self._tf_update_params(w_mu, w_Sigma, w_Lambda)


  def resample_w(self, cholesky=False):
    sample = tf.random_normal(shape=self.w_mu.shape, dtype=self.dtype)

    # Compute A s.t. A A^T = w_Sigma. Note that SVD and Cholesky give different A
    if cholesky:
      # Use cholesky
      A = tf.cholesky(self.w_Sigma)
    else:
      # Use SVD
      S, U, _ = tf.svd(self.w_Sigma)
      A = tf.matmul(U, tf.diag(tf.sqrt(S)))

    w = self.w_mu + tf.matmul(A, sample)
    return tf.assign(self.w, w, name="resample_w")


  @property
  def reset(self):
    return self.reset_op


  @property
  def trainable_weights(self):
    return self._trainable_weights or []


  def _tf_update_params(self, w_mu, w_Sigma, w_Lambda):
    """
    Returns:
      tf.Op which performs an update on all weight parameters
    """
    mu_op     = tf.assign(self.w_mu,      w_mu)
    Sigma_op  = tf.assign(self.w_Sigma,   w_Sigma)
    Lambda_op = tf.assign(self.w_Lambda,  w_Lambda)
    return tf.group(mu_op, Sigma_op, Lambda_op)


  def _cast_input(self, x):
    if self.dtype == tf.float64 and x.dtype.base_dtype != tf.float64:
      x = tf.cast(x, self.dtype)
    return x


  def _cast_output(self, x):
    if x.dtype.base_dtype != tf.float32:
      x = tf.cast(x, tf.float32)
    return x
