# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import ast
import numpy as np
import mxnet as mx

class BinaryRBM(mx.operator.CustomOp):

    def __init__(self, k):
        self.k = k # Persistent contrastive divergence k

    def forward(self, is_train, req, in_data, out_data, aux):
        visible_layer_data = in_data[0] # (num_batch, num_visible)
        visible_layer_bias = in_data[1] # (num_visible,)
        hidden_layer_bias = in_data[2]  # (num_hidden,)
        interaction_weight= in_data[3]        # (num_visible, num_hidden)

        if is_train:
            _, hidden_layer_prob_1 = self.sample_hidden_layer(visible_layer_data, hidden_layer_bias, interaction_weight)
            hidden_layer_sample = aux[1] # The initial state of the Gibbs sampling for persistent CD
        else:
            hidden_layer_sample, hidden_layer_prob_1 = self.sample_hidden_layer(visible_layer_data, hidden_layer_bias, interaction_weight)

        # k-step Gibbs sampling
        for _ in range(self.k):
            visible_layer_sample, visible_layer_prob_1 = self.sample_visible_layer(hidden_layer_sample, visible_layer_bias, interaction_weight)
            hidden_layer_sample, _ = self.sample_hidden_layer(visible_layer_sample, hidden_layer_bias, interaction_weight)

        if is_train:
            # Used in backward and next forward
            aux[0][:] = visible_layer_sample
            aux[1][:] = hidden_layer_sample

        self.assign(out_data[0], req[0], visible_layer_prob_1)
        self.assign(out_data[1], req[1], hidden_layer_prob_1)

    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        visible_layer_data = in_data[0]    # (num_batch, num_visible)
        visible_layer_sample = aux[0]      # (num_batch, num_visible)
        hidden_layer_prob_1 = out_data[1]  # (num_batch, num_hidden)
        hidden_layer_sample = aux[1]       # (num_batch, num_hidden)

        grad_visible_layer_bias = (visible_layer_sample - visible_layer_data).mean(axis=0)
        grad_hidden_layer_bias = (hidden_layer_sample - hidden_layer_prob_1).mean(axis=0)
        grad_interaction_weight= (mx.nd.linalg.gemm2(visible_layer_sample.expand_dims(2), hidden_layer_sample.expand_dims(1)) -
                            mx.nd.linalg.gemm2(visible_layer_data.expand_dims(2), hidden_layer_prob_1.expand_dims(1))
                           ).mean(axis=0)

        # We don't need the gradient on the visible layer input
        self.assign(in_grad[1], req[1], grad_visible_layer_bias)
        self.assign(in_grad[2], req[2], grad_hidden_layer_bias)
        self.assign(in_grad[3], req[3], grad_interaction_weight)

    def sample_hidden_layer(self, visible_layer_batch, hidden_layer_bias, interaction_weight):
        return self.sample_layer(visible_layer_batch, hidden_layer_bias, interaction_weight, False)

    def sample_visible_layer(self, hidden_layer_batch, visible_layer_bias, interaction_weight):
        return self.sample_layer(hidden_layer_batch, visible_layer_bias, interaction_weight, True)

    def sample_layer(self, other_layer_sample, layer_bias, interaction_weight, interaction_transpose):
        prob_1 = mx.nd.linalg.gemm(
            other_layer_sample,
            interaction_weight,
            layer_bias.tile(reps=(other_layer_sample.shape[0], 1)),
            transpose_b=interaction_transpose) # (num_batch, num_units_in_layer)
        prob_1.sigmoid(out=prob_1)
        return mx.nd.random.uniform(shape=prob_1.shape) < prob_1, prob_1

@mx.operator.register('BinaryRBM')
class BinaryRBMProp(mx.operator.CustomOpProp):

    # Auxiliary states are requested only if `for_training` is true.
    def __init__(self, num_hidden, k, for_training):
        super(BinaryRBMProp, self).__init__(False)
        self.num_hidden = int(num_hidden)
        self.k = int(k)
        self.for_training = ast.literal_eval(for_training)

    def list_arguments(self):
        # 0: (batch size, the number of visible units)
        # 1: (the number of visible units,)
        # 2: (the number of hidden units,)
        # 3: (the number of visible units, the number of hidden units)
        return ['data', 'visible_layer_bias', 'hidden_layer_bias', 'interaction_weight']

    def list_outputs(self):
        # 0: The probabilities that each visible unit is 1 after `k` steps of Gibbs sampling starting from the given `data`.
        #    (batch size, the number of visible units)
        # 1: The probabilities that each hidden unit is 1 conditional on the given `data`.
        #    (batch size, the number of hidden units)
        return ['visible_layer_prob_1', 'hidden_layer_prob_1']

    def list_auxiliary_states(self):
        # Used only if `self.for_trainig is true.
        # 0: Store the visible layer samples obtained in the forward pass, used in the backward pass.
        #    (batch size, the number of visible units)
        # 1: Store the hidden layer samples obtained in the forward pass, used in the backward and next forward pass.
        #    (batch size, the number of hidden units)
        return ['aux_visible_layer_sample', 'aux_hidden_layer_sample'] if self.for_training else []

    def infer_shape(self, in_shapes):
        visible_layer_data_shape = in_shapes[0] # The input data
        visible_layer_bias_shape = (visible_layer_data_shape[1],)
        hidden_layer_bias_shape = (self.num_hidden,)
        interaction_shape = (visible_layer_data_shape[1], self.num_hidden)
        visible_layer_sample_shape = visible_layer_data_shape
        visible_layer_prob_1_shape = visible_layer_sample_shape
        hidden_layer_sample_shape = (visible_layer_data_shape[0], self.num_hidden)
        hidden_layer_prob_1_shape = hidden_layer_sample_shape
        return [visible_layer_data_shape, visible_layer_bias_shape, hidden_layer_bias_shape, interaction_shape], \
               [visible_layer_prob_1_shape, hidden_layer_prob_1_shape], \
               [visible_layer_sample_shape, hidden_layer_sample_shape] if self.for_training else []

    def infer_type(self, in_type):
        return [in_type[0], in_type[0], in_type[0], in_type[0]], \
               [in_type[0], in_type[0]], \
               [in_type[0], in_type[0]] if self.for_training else []

    def create_operator(self, ctx, in_shapes, in_dtypes):
        return BinaryRBM(self.k)

# For gluon API
class BinaryRBMBlock(mx.gluon.HybridBlock):

    def __init__(self, num_hidden, k, for_training, **kwargs):
        super(BinaryRBMBlock, self).__init__(**kwargs)
        with self.name_scope():
            self.num_hidden = num_hidden
            self.k = k
            self.for_training = for_training
            self.visible_layer_bias = self.params.get('visible_layer_bias', shape=(0,), allow_deferred_init=True)
            self.hidden_layer_bias = self.params.get('hidden_layer_bias', shape=(0,), allow_deferred_init=True)
            self.interaction_weight= self.params.get('interaction_weight', shape=(0, 0), allow_deferred_init=True)
            if for_training:
                self.aux_visible_layer_sample = self.params.get('aux_visible_layer_sample', shape=(0, 0), allow_deferred_init=True)
                self.aux_hidden_layer_sample = self.params.get('aux_hidden_layer_sample', shape=(0, 0), allow_deferred_init=True)

    def hybrid_forward(self, F, data, visible_layer_bias, hidden_layer_bias, interaction_weight, aux_visible_layer_sample=None, aux_hidden_layer_sample=None):
        # As long as `for_training` is kept constant, this conditional statement does not prevent hybridization.
        if self.for_training:
            return F.Custom(
                data,
                visible_layer_bias,
                hidden_layer_bias,
                interaction_weight,
                aux_visible_layer_sample,
                aux_hidden_layer_sample,
                num_hidden=self.num_hidden,
                k=self.k,
                for_training=self.for_training,
                op_type='BinaryRBM')
        else:
            return F.Custom(
                data,
                visible_layer_bias,
                hidden_layer_bias,
                interaction_weight,
                num_hidden=self.num_hidden,
                k=self.k,
                for_training=self.for_training,
                op_type='BinaryRBM')

def estimate_log_likelihood(visible_layer_bias, hidden_layer_bias, interaction_weight, ais_batch_size, ais_num_batch, ais_intermediate_steps, ais_burn_in_steps, data, ctx):
    # The base-rate RBM with no hidden layer. The visible layer bias is set to the same with the given RBM.
    # This is not the only possible choice but simple and works well.
    base_rate_visible_layer_bias = visible_layer_bias
    base_rate_visible_prob_1 = base_rate_visible_layer_bias.sigmoid()
    log_base_rate_z = base_rate_visible_layer_bias.exp().log1p().sum()

    def log_intermediate_unnormalized_prob(visible_layer_sample, beta):
        p = mx.nd.dot(
                visible_layer_sample, 
                (1 - beta) * base_rate_visible_layer_bias + beta * visible_layer_bias)
        if beta != 0:
            p += mx.nd.linalg.gemm(
                    visible_layer_sample,
                    interaction_weight,
                    hidden_layer_bias.tile(reps=(visible_layer_sample.shape[0], 1)),
                    transpose_b=False,
                    alpha=beta,
                    beta=beta).exp().log1p().sum(axis=1)
        return p

    def sample_base_rbm():
        rands = mx.nd.random.uniform(shape=(ais_batch_size, base_rate_visible_prob_1.shape[0]), ctx=ctx)
        return rands < base_rate_visible_prob_1.tile(reps=(ais_batch_size, 1))

    def sample_intermediate_visible_layer(visible_layer_sample, beta):
        for _ in range(ais_burn_in_steps):
            hidden_prob_1 = mx.nd.linalg.gemm(
                visible_layer_sample,
                interaction_weight,
                hidden_layer_bias.tile(reps=(visible_layer_sample.shape[0], 1)),
                transpose_b=False,
                alpha=beta,
                beta=beta)
            hidden_prob_1.sigmoid(out=hidden_prob_1)
            hidden_layer_sample = mx.nd.random.uniform(shape=hidden_prob_1.shape, ctx=ctx) < hidden_prob_1
            visible_prob_1 = mx.nd.linalg.gemm(
                hidden_layer_sample,
                interaction_weight,
                visible_layer_bias.tile(reps=(hidden_layer_sample.shape[0], 1)),
                transpose_b=True,
                alpha=beta,
                beta=beta) + (1 - beta) * base_rate_visible_layer_bias
            visible_prob_1.sigmoid(out=visible_prob_1)
            visible_layer_sample = mx.nd.random.uniform(shape=visible_prob_1.shape, ctx=ctx) < visible_prob_1
        return visible_layer_sample

    def array_from_batch(batch):
        if isinstance(batch, mx.io.DataBatch):
            return batch.data[0].as_in_context(ctx).flatten()
        else: # batch is an instance of list in the case of gluon DataLoader
            return batch[0].as_in_context(ctx).flatten()

    importance_weight_sum = 0
    num_ais_samples = ais_num_batch * ais_batch_size
    for _ in range(ais_num_batch):
        log_importance_weight = 0
        visible_layer_sample = sample_base_rbm()
        for n in range(1, ais_intermediate_steps + 1):
            beta = 1. * n / ais_intermediate_steps
            log_importance_weight += \
                log_intermediate_unnormalized_prob(visible_layer_sample, beta) - \
                log_intermediate_unnormalized_prob(visible_layer_sample, (n - 1.) / ais_intermediate_steps)
            visible_layer_sample = sample_intermediate_visible_layer(visible_layer_sample, beta)
        importance_weight_sum += log_importance_weight.exp().sum()
    log_z = (importance_weight_sum / num_ais_samples).log() + log_base_rate_z

    log_likelihood = 0
    num_data = 0
    for batch in data:
        batch_array = array_from_batch(batch)
        log_likelihood += log_intermediate_unnormalized_prob(batch_array, 1) - log_z
        num_data += batch_array.shape[0]
    log_likelihood = log_likelihood.sum() / num_data

    return log_likelihood.asscalar(), log_z.asscalar()