"""Memory module for Kanerva Machines.

  Functions of the module always take inputs with shape:
  [seq_length, batch_size, ...]

  Examples:

    # Initialisation
    memory = KanervaMemory(code_size=100, memory_size=32)
    prior_memory = memory.get_prior_state(batch_size)

    # Update memory posterior
    posterior_memory, _, _, _ = memory.update_state(z_episode, prior_memory)

    # Read from the memory using cues z_q
    read_z, dkl_w = memory.read_with_z(z_q, posterior_memory)

    # Compute the KL-divergence between posterior and prior memory
    dkl_M = memory.get_dkl_total(posterior_memory)
"""
import collections
import functools

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.python.keras.layers import Layer
from tensorflow.python.keras.initializers import constant, zeros, truncated_normal


def define_scope(scope: str = None):
    def decorator(function):
        name_scope = scope or function.__name__

        @functools.wraps(function)
        def decorated(*args, **kwargs):
            with tf.name_scope(name_scope):
                return function(*args, **kwargs)

        return decorated

    return decorator


MemoryState = collections.namedtuple(
    'MemoryState',
    # Mean of memory slots, [batch_size, memory_size, word_size]
    # Covariance of memory slots, [batch_size, memory_size, memory_size]
    ('M_mean', 'M_cov'))

EPSILON = 1e-6


# disable lint warnings for cleaner algebraic expressions
# pylint: disable=invalid-name
class KanervaMemory(Layer):
    """A memory-based generative model."""

    def __init__(self,
                 code_size,
                 memory_size,
                 num_opt_iters=1,
                 w_prior_stddev=1.0,
                 obs_noise_stddev=1.0,
                 sample_w=False,
                 sample_M=False,
                 name='KanervaMemory'):
        """Initialise the memory module.

        Args:
          code_size: Integer specifying the size of each encoded input.
          memory_size: Integer specifying the total number of rows in the memory.
          num_opt_iters: Integer specifying the number of optimisation iterations.
          w_prior_stddev: Float specifying the standard deviation of w's prior.
          obs_noise_stddev: Float specifying the standard deviation of the
            observational noise.
          sample_w: Boolean specifying whether to sample w or simply take its mean.
          sample_M: Boolean specifying whether to sample M or simply take its mean.
          name: String specfying the name of this module.
        """
        super().__init__(name=name)
        self._memory_size = memory_size
        self._code_size = code_size
        self._num_opt_iters = num_opt_iters
        self._sample_w = sample_w
        self._sample_M = sample_M
        self._w_prior_stddev = tf.constant(w_prior_stddev)

        with tf.name_scope(name):
            log_w_stddev = tf.constant(np.log(0.3))
            if obs_noise_stddev > 0.0:
                self._obs_noise_stddev = tf.constant(obs_noise_stddev)
            else:
                log_obs_stddev = tf.constant(0)
                self._obs_noise_stddev = tf.exp(log_obs_stddev)
        self._w_stddev = tf.exp(log_w_stddev)
        self._w_prior_dist = tfp.distributions.MultivariateNormalDiag(
            loc=tf.zeros([self._memory_size]),
            scale_identity_multiplier=self._w_prior_stddev)

    def _build(self):
        raise ValueError('`_build()` should not be called for this module since'
                         'it takes no inputs and all of its variables are'
                         'constructed in `__init__`')

    def _get_w_dist(self, mu_w):
        return tfp.distributions.MultivariateNormalDiag(
            loc=mu_w, scale_identity_multiplier=self._w_stddev)

    def sample_prior_w(self, seq_length, batch_size):
        """Sample w from its prior.

        Args:
          seq_length: length of sequence
          batch_size: batch size of samples
        Returns:
          w: [batch_size, memory_size]
        """
        return self._w_prior_dist.sample([seq_length, batch_size])

    def read_with_z(self, z, memory_state):
        """Query from memory (specified by memory_state) using embedding z.

        Args:
          z: Tensor with dimensions [episode_length, batch_size, code_size]
            containing an embedded input.
          memory_state: Instance of `MemoryState`.

        Returns:
          A tuple of tensors containing the mean of read embedding and the
            KL-divergence between the w used in reading and its prior.
        """
        M = self.sample_M(memory_state)
        w_mean = self._solve_w_mean(z, M)
        w_samples = self.sample_w(w_mean)
        dkl_w = self.get_dkl_w(w_mean)
        z_mean = self.get_w_to_z_mean(w_samples, M)
        return z_mean, dkl_w

    def wrap_z_dist(self, z_mean):
        """Wrap the mean of z as an observation (Gaussian) distribution."""
        return tfp.distributions.MultivariateNormalDiag(
            loc=z_mean, scale_identity_multiplier=self._obs_noise_stddev)

    def sample_w(self, w_mean):
        """Sample w from its posterior distribution."""
        if self._sample_w:
            return self._get_w_dist(w_mean).sample()
        else:
            return w_mean

    def sample_M(self, memory_state):
        """Sample the memory from its distribution specified by memory_state."""
        if self._sample_M:
            noise_dist = tfp.distributions.MultivariateNormalFullCovariance(covariance_matrix=memory_state.M_cov)
            # C, B, M
            noise = tf.transpose(noise_dist.sample(self._code_size),
                                 [1, 2, 0])
            return memory_state.M_mean + noise
        else:
            return memory_state.M_mean

    def get_w_to_z_mean(self, w_p, R):
        """Return the mean of z by reading from memory using weights w_p."""
        return tf.einsum('sbm,bmc->sbc', w_p, R)  # Rw

    def _read_cov(self, w_samples, memory_state):
        episode_size, batch_size = w_samples.get_shape().as_list()[:2]
        _, U = memory_state  # cov: [B, M, M]
        wU = tf.einsum('sbm,bmn->sbn', w_samples, U)
        wUw = tf.einsum('sbm,sbm->sb', wU, w_samples)
        wUw.get_shape().assert_is_compatible_with([episode_size, batch_size])
        return wU, wUw

    def get_dkl_total(self, memory_state):
        """Compute the KL-divergence between a memory distribution and its prior."""
        R, U = memory_state
        B, K, _ = R.get_shape().as_list()
        U.get_shape().assert_is_compatible_with([B, K, K])
        R_prior, U_prior = self.get_prior_state(B)
        p_diag = tf.matrix_diag_part(U_prior)
        q_diag = tf.matrix_diag_part(U)  # B, K
        t1 = self._code_size * tf.reduce_sum(q_diag / p_diag, -1)
        t2 = tf.reduce_sum((R - R_prior) ** 2 / tf.expand_dims(
            p_diag, -1), [-2, -1])
        t3 = -self._code_size * self._memory_size
        t4 = self._code_size * tf.reduce_sum(tf.log(p_diag) - tf.log(q_diag), -1)
        return t1 + t2 + t3 + t4

    def _get_dkl_update(self, memory_state, w_samples, new_z_mean, new_z_var):
        """Compute memory_kl after updating prior_state."""
        B, K, C = memory_state.M_mean.get_shape().as_list()
        S = w_samples.get_shape().as_list()[0]

        # check shapes
        w_samples.get_shape().assert_is_compatible_with([S, B, K])
        new_z_mean.get_shape().assert_is_compatible_with([S, B, C])

        delta = new_z_mean - self.get_w_to_z_mean(w_samples, memory_state.M_mean)
        _, wUw = self._read_cov(w_samples, memory_state)
        var_z = wUw + new_z_var + self._obs_noise_stddev ** 2
        beta = wUw / var_z

        dkl_M = -0.5 * (self._code_size * beta
                        - tf.reduce_sum(tf.expand_dims(beta / var_z, -1)
                                        * delta ** 2, -1)
                        + self._code_size * tf.log(1 - beta))
        dkl_M.get_shape().assert_is_compatible_with([S, B])
        return dkl_M

    @define_scope()
    def _get_prior_params(self):
        log_var = tf.Variable(0., name='prior_var_scale')
        self._prior_var = tf.ones([self._memory_size]) * tf.exp(log_var) + EPSILON
        prior_cov = tf.matrix_diag(self._prior_var)
        prior_mean = tf.Variable(
            tf.truncated_normal_initializer(mean=0.0, stddev=1.0)([self._memory_size, self._code_size]),
            name='prior_mean'
        )
        return prior_mean, prior_cov

    @property
    def prior_avg_var(self):
        """return the average of prior memory variance."""
        return tf.reduce_mean(self._prior_var)

    def _solve_w_mean(self, new_z_mean, M):
        """Minimise the conditional KL-divergence between z wrt w."""
        w_matrix = tf.matmul(M, M, transpose_b=True)
        w_rhs = tf.einsum('bmc,sbc->bms', M, new_z_mean)
        w_mean = tf.matrix_solve_ls(
            matrix=w_matrix, rhs=w_rhs,
            l2_regularizer=self._obs_noise_stddev ** 2 / self._w_prior_stddev ** 2)
        w_mean = tf.einsum('bms->sbm', w_mean)
        return w_mean

    def get_prior_state(self, batch_size):
        """Return the prior distribution of memory as a MemoryState."""
        prior_mean, prior_cov = self._get_prior_params()
        batch_prior_mean = tf.stack([prior_mean] * batch_size)
        batch_prior_cov = tf.stack([prior_cov] * batch_size)
        return MemoryState(M_mean=batch_prior_mean,
                           M_cov=batch_prior_cov)

    def update_state(self, z, memory_state):
        """Update the memory state using Bayes' rule.

        Args:
          z: A tensor with dimensions [episode_length, batch_size, code_size]
            containing a sequence of embeddings to write into memory.
          memory_state: A `MemoryState` namedtuple containing the memory state to
            be written to.

        Returns:
          A tuple containing the following elements:
          final_memory: A `MemoryState` namedtuple containing the new memory state
            after the update.
          w_mean_episode: The mean of w for the written episode.
          dkl_w_episode: The KL-divergence of w for the written episode.
          dkl_M_episode: The KL-divergence between the memory states before and
            after the update.
        """

        episode_size, batch_size = z.get_shape().as_list()[:2]
        w_array = tf.TensorArray(dtype=tf.float32, size=episode_size,
                                 element_shape=[1, batch_size, self._memory_size])
        dkl_w_array = tf.TensorArray(dtype=tf.float32, size=episode_size,
                                     element_shape=[1, batch_size])
        dkl_M_array = tf.TensorArray(dtype=tf.float32, size=episode_size,
                                     element_shape=[1, batch_size])
        init_var = (0, memory_state, w_array, dkl_w_array, dkl_M_array)
        cond = lambda i, m, d_2, d_3, d_4: i < episode_size

        def loop_body(i, old_memory, w_array, dkl_w_array, dkl_M_array):
            """Update memory step-by-step."""
            z_step = tf.expand_dims(z[i], 0)
            new_memory = old_memory
            for _ in range(self._num_opt_iters):
                w_step_mean = self._solve_w_mean(z_step, self.sample_M(new_memory))
                w_step_sample = self.sample_w(w_step_mean)
                new_memory = self._update_memory(old_memory,
                                                 w_step_mean,
                                                 z_step, 0)
            dkl_w_step = self.get_dkl_w(w_step_mean)
            dkl_M_step = self._get_dkl_update(old_memory,
                                              w_step_sample,
                                              z_step, 0)
            return (i + 1,
                    new_memory,
                    w_array.write(i, w_step_sample),
                    dkl_w_array.write(i, dkl_w_step),
                    dkl_M_array.write(i, dkl_M_step))

        _, final_memory, w_mean, dkl_w, dkl_M = tf.while_loop(
            cond, loop_body, init_var)
        w_mean_episode = w_mean.concat()
        dkl_w_episode = dkl_w.concat()
        dkl_M_episode = dkl_M.concat()
        dkl_M_episode.get_shape().assert_is_compatible_with(
            [episode_size, batch_size])

        return final_memory, w_mean_episode, dkl_w_episode, dkl_M_episode

    def _update_memory(self, old_memory, w_samples, new_z_mean, new_z_var):
        """Setting new_z_var=0 for sample based update."""
        old_mean, old_cov = old_memory
        wR = self.get_w_to_z_mean(w_samples, old_memory.M_mean)
        wU, wUw = self._read_cov(w_samples, old_memory)
        sigma_z = wUw + new_z_var + self._obs_noise_stddev ** 2  # [S, B]
        delta = new_z_mean - wR  # [S, B, C]
        c_z = wU / tf.expand_dims(sigma_z, -1)  # [S, B, M]
        posterior_mean = old_mean + tf.einsum('sbm,sbc->bmc', c_z, delta)
        posterior_cov = old_cov - tf.einsum('sbm,sbn->bmn', c_z, wU)
        # Clip diagonal elements for numerical stability
        posterior_cov = tf.matrix_set_diag(
            posterior_cov,
            tf.clip_by_value(tf.matrix_diag_part(posterior_cov), EPSILON, 1e10))
        new_memory = MemoryState(M_mean=posterior_mean, M_cov=posterior_cov)
        return new_memory

    def get_dkl_w(self, w_mean):
        """Return the KL-divergence between posterior and prior weights w."""
        posterior_dist = self._get_w_dist(w_mean)
        dkl_w = posterior_dist.kl_divergence(self._w_prior_dist)
        dkl_w.get_shape().assert_is_compatible_with(
            w_mean.get_shape().as_list()[:-1])
        return dkl_w