# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import numpy as np

DebugOutput = False
np.set_printoptions(suppress=True) #, precision=16, floatmode='maxprec')

def print_with_shape(name, a, force_output=False):
    if (force_output or DebugOutput):
        print(name + " [shape: ", a.shape, "]\n", a)

def print_results(Y):
    print("*************************")
    print_with_shape("Y", Y, True)
    print("*************************")

class GRU_Helper():

    def __init__(self, **params):
        # Match the ONNXRuntime/CNTK behavior
        # If False use the python from the ONNX spec
        self.match_onnxruntime = True

        required_inputs = ['X', 'W', 'R']
        for i in required_inputs:
            assert i in params, "Missing Required Input: {0}".format(i)

        num_directions = params['W'].shape[0]
        sequence_length = params['X'].shape[0]

        hidden_size = params['R'].shape[-1]
        batch_size = params['X'].shape[1]

        X = params['X']
        W = params['W']
        R = params['R']
        B = params['B'] if 'B' in params else np.zeros(num_directions * 6 * hidden_size).reshape(num_directions, 6 * hidden_size)
        H_0 = params['initial_h'] if 'initial_h' in params else np.zeros((num_directions, batch_size, hidden_size)).reshape(num_directions, batch_size, hidden_size)
        LBR = params['linear_before_reset'] if 'linear_before_reset' in params else 0
        self.direction = params['direction'] if 'direction' in params else 'forward'

        if (num_directions == 1):
            if (self.direction == 'forward'):
                self.one = OneDirectionGRU(X, W, R, B, H_0, LBR)
            else:
                # flip input so we process in reverse
                self.one = OneDirectionGRU(np.flip(X,0), W, R, B, H_0, LBR)

            self.two = None

        else:
            # split the inputs which have per direction rows
            Wfw, Wbw = np.vsplit(W, 2)
            Rfw, Rbw = np.vsplit(R, 2)
            Bfw, Bbw = np.vsplit(B, 2)
            H_0fw, H_0bw = np.vsplit(H_0, 2)

            self.one = OneDirectionGRU(X, Wfw, Rfw, Bfw, H_0fw, LBR)
            self.two = OneDirectionGRU(np.flip(X, 0), Wbw, Rbw, Bbw, H_0bw, LBR)

    def run(self):

        if (self.direction == 'bidirectional'):
            f_output = self.one.execute()
            r_output = self.two.execute()

            # flip reverse output it matches the original input order
            r_output_orig_input_order = np.flip(r_output, 0)

            # create merged output by merging the forward and reverse rows for seq_length
            # 0 rows, 2 directions, batch size, hidden_size
            seq_length = f_output.shape[0]
            batch_size = f_output.shape[2]
            hidden_size = f_output.shape[3]

            output = np.empty((0, 2, batch_size, hidden_size), np.float32)
            for x in range(0, seq_length):
                output = np.append(output, f_output[x])
                output = np.append(output, r_output_orig_input_order[x])

            output = output.reshape(seq_length, 2, batch_size, hidden_size)
        else:
            output = self.one.execute()
            if (self.direction == 'reverse'):
                # flip so it's back in the original order of the inputs
                output = np.flip(output, 0)

        return output

class OneDirectionGRU():

    def __init__(self, X, W, R, B, initial_h, LBR):

        self.X = X
        # remove num_directions axis for W, R, B, H_0
        self.W = np.squeeze(W, axis=0)
        self.R = np.squeeze(R, axis=0)
        self.B = np.squeeze(B, axis=0)
        self.H_0 = np.squeeze(initial_h, axis=0)
        self.LBR = LBR

    def f(self, x):
        return 1 / (1 + np.exp(-x))

    def g(self, x):
        return np.tanh(x)

    def execute(self):
        print_with_shape("X", self.X)

        [w_z, w_r, w_h] = np.split(self.W, 3)
        [r_z, r_r, r_h] = np.split(self.R, 3)
        [w_bz, w_br, w_bh, r_bz, r_br, r_bh] = np.split(self.B, 6)

        #print_with_shape("w_z", w_z)
        #print_with_shape("w_r", w_r)
        #print_with_shape("w_h", w_h)

        #print_with_shape("r_z", r_z)
        #print_with_shape("r_r", r_r)
        #print_with_shape("r_h", r_h)

        #print_with_shape("w_bz", w_bz)
        #print_with_shape("w_br", w_br)
        #print_with_shape("w_bh", w_bh)
        #print_with_shape("r_bz", r_bz)
        #print_with_shape("r_br", r_br)
        #print_with_shape("r_bh", r_bh)

        seq_len = self.X.shape[0]
        num_directions = 1
        hidden_size = self.R.shape[-1]
        batch_size = self.X.shape[1]

        output = np.empty((0, num_directions, batch_size, hidden_size), np.float32)

        for row in self.X:
            z = self.f(np.dot(row, np.transpose(w_z)) + np.dot(self.H_0, np.transpose(r_z)) + w_bz + r_bz)
            r = self.f(np.dot(row, np.transpose(w_r)) + np.dot(self.H_0, np.transpose(r_r)) + w_br + r_br)
            h_default = self.g(np.dot(row, np.transpose(w_h)) + np.dot(r * self.H_0, np.transpose(r_h)) + w_bh + r_bh)
            h_linear = self.g(np.dot(row, np.transpose(w_h)) + r * (np.dot(self.H_0, np.transpose(r_h)) + r_bh) + w_bh)

            h = h_linear if self.LBR else h_default

            print_with_shape("z", z)
            print_with_shape("r", r)
            print_with_shape("h", h)

            H = (1 - z) * h + z * self.H_0

            print_with_shape("H", H)
            output = np.append(output, H.reshape(1, 1, batch_size, hidden_size), axis=0)

            self.H_0 = H

        return output

class ONNXRuntimeTestContext():

    @staticmethod
    def OneDirectionWeights():

        hidden_size = 2

        W = np.array([[[-0.494659, 0.0453352],  # Wz
                       [-0.487793, 0.417264],
                       [-0.0091708, -0.255364],  # Wr
                       [-0.106952, -0.266717],
                       [-0.0888852, -0.428709],  # Wh
                       [-0.283349, 0.208792]]]).astype(np.float32)

        R = np.array([[[0.146626, -0.0620289],  # Rz
                       [-0.0815302, 0.100482],
                       [-0.228172, 0.405972],  # Rr
                       [0.31576, 0.281487],
                       [-0.394864, 0.42111],  # Rh
                       [-0.386624, -0.390225]]]).astype(np.float32)

        W_B = np.array([[0.381619, 0.0323954,  # Wbz
                         -0.258721, 0.45056,  # Wbr
                         -0.250755, 0.0967895]]).astype(np.float32)  # Wbh
        R_B = np.zeros((1, 3 * hidden_size)).astype(np.float32)
        B = np.concatenate((W_B, R_B), axis=1)

        return W, R, B

    @staticmethod
    def BidirectionalWeights():

        W1, R1, B1 = ONNXRuntimeTestContext.OneDirectionWeights()

        hidden_size = R1.shape[-1]
        input_size = W1.shape[-1]

        W = np.tile(W1, (2, 1)).reshape(2, 3 * hidden_size, input_size)
        R = np.tile(R1, (2, 1)).reshape(2, 3 * hidden_size, hidden_size)
        B = np.tile(B1, (2,1))

        return W, R, B

# replicate ONNXRuntime unit tests inputs to validate output
class GRU_ONNXRuntimeUnitTests():

    @staticmethod
    def ForwardDefaultActivationsSimpleWeightsNoBiasTwoRows():

        print(GRU_ONNXRuntimeUnitTests.ForwardDefaultActivationsSimpleWeightsNoBiasTwoRows.__name__)

        seq_length = 2
        batch_size = 2
        input_size = 1
        hidden_size = 3
        input = np.array([1., 2., 10., 11.]).astype(np.float32).reshape(seq_length, batch_size, input_size)

        W = np.array([0.1, 0.2, 0.3, 1, 2, 3, 10, 11, 12]).astype(np.float32).reshape(1, 3 * hidden_size, input_size)

        weight_scale = 0.1
        R = weight_scale * np.ones((1, 3 * hidden_size, hidden_size)).astype(np.float32)

        gru = GRU_Helper(X=input,W=W,R=R,direction='forward')
        fw_output = gru.run()
        print_results(fw_output)

    @staticmethod
    def ReverseDefaultActivationsSimpleWeightsNoBiasTwoRows():

        print(GRU_ONNXRuntimeUnitTests.ReverseDefaultActivationsSimpleWeightsNoBiasTwoRows.__name__)

        seq_length = 2
        batch_size = 2
        input_size = 1
        hidden_size = 3
        input = np.array([[[1.], [2.]],
                          [[10.], [11.]]]).astype(np.float32)


        W = np.array([0.1, 0.2, 0.3, 1, 2, 3, 10, 11, 12]).astype(np.float32).reshape(1, 3 * hidden_size, input_size)

        weight_scale = 0.1
        R = weight_scale * np.ones((1, 3 * hidden_size, hidden_size)).astype(np.float32)

        gru = GRU_Helper(X=input,W=W,R=R,direction='reverse')
        fw_output = gru.run()
        print_results(fw_output)

    @staticmethod
    def BidirectionalDefaultActivationsSimpleWeightsNoBiasTwoRows():

        print(GRU_ONNXRuntimeUnitTests.BidirectionalDefaultActivationsSimpleWeightsNoBiasTwoRows.__name__)

        seq_length = 2
        batch_size = 2
        input_size = 1
        hidden_size = 3

        input = np.array([[[1.], [2.]],
                          [[10.], [11.]]]).astype(np.float32)

        W = np.array([0.1, 0.2, 0.3, 1, 2, 3, 10, 11, 12]).astype(np.float32).reshape(1, 3 * hidden_size, input_size)

        weight_scale = 0.1
        R = weight_scale * np.ones((1, 3 * hidden_size, hidden_size)).astype(np.float32)

        # duplicate the W and R inputs so we use the same values for both forward and reverse
        gru = GRU_Helper(X=input,
                         W=np.tile(W,(2,1)).reshape(2, 3 * hidden_size, input_size),
                         R=np.tile(R,(2,1)).reshape(2, 3 * hidden_size, hidden_size),
                         direction='bidirectional')

        fw_output = gru.run()
        print_results(fw_output)

    @staticmethod
    def DefaultActivationsSimpleWeightsWithBias(rows=2, direction="forward", linear_before_reset=0):

        print(GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias.__name__ +
              " batch_parallel=" + str(rows != 1) +
              " direction=" + direction +
              " linear_before_reset=" + str(linear_before_reset))

        seq_length = 2
        batch_size = rows
        input_size = 1
        hidden_size = 3

        if (batch_size == 1):
            input = [-0.1, -0.3]
        else:
            input = [-0.1, 0.2, -0.3, 0.4]

        input = np.array(input).astype(np.float32).reshape(seq_length, batch_size, input_size)

        W = np.array([0.1, 0.2, 0.3, 0.2, 0.3, 0.1, 0.3, 0.1, 0.2]).astype(np.float32).reshape(1, 3 * hidden_size, input_size)

        weight_scale = 0.1
        R = weight_scale * np.ones((1, 3 * hidden_size, hidden_size)).astype(np.float32)

        # Wb[zrh] Rb[zrh]
        B = np.array([-0.01, 0.1, 0.01, -0.2, -0.02, 0.02, 0.3, -0.3, -0.3,
                      -0.03, 0.5, -0.7, 0.05, -0.7, 0.3, 0.07, -0.03, 0.5 ]).astype(np.float32).reshape(1, 6 * hidden_size)

        gru = GRU_Helper(X=input,W=W,R=R, B=B, direction=direction, linear_before_reset=linear_before_reset)
        fw_output = gru.run()
        print_results(fw_output)

    @staticmethod
    def ForwardDefaultActivationsSimpleWeightsWithBiasBatchParallel():
        GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias()

    @staticmethod
    def ForwardDefaultActivationsSimpleWeightsWithBiasBatchParallelLinearBeforeReset():

        GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias(linear_before_reset=1)

    @staticmethod
    def ReverseDefaultActivationsSimpleWeightsWithBiasBatchParallelLinearBeforeReset(linear_before_reset=0):

        GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias(direction="reverse", linear_before_reset=1)

    @staticmethod
    def ForwardDefaultActivationsSimpleWeightsWithBiasLinearBeforeReset():

        GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias(rows=1, linear_before_reset=1)

    @staticmethod
    def ReverseDefaultActivationsSimpleWeightsWithBiasLinearBeforeReset(linear_before_reset=0):

        GRU_ONNXRuntimeUnitTests.DefaultActivationsSimpleWeightsWithBias(rows=1, direction="reverse", linear_before_reset=1)

    @staticmethod
    def Legacy_TestGRUOpForwardBasic():

        print(GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpForwardBasic.__name__)

        input = np.array([[[-0.455351, -0.276391]],
                          [[-0.185934, -0.269585]]]).astype(np.float32)

        W, R, B = ONNXRuntimeTestContext.OneDirectionWeights()
        gru = GRU_Helper(X=input, W=W, R=R, B=B)
        output = gru.run()
        print_results(output)

    @staticmethod
    def Legacy_TestGRUOpBackwardBasic():
        print(GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpBackwardBasic.__name__)

        input = np.array([[[-0.185934, -0.269585]],
                          [[-0.455351, -0.276391]]]).astype(np.float32)

        W, R, B = ONNXRuntimeTestContext.OneDirectionWeights()
        gru = GRU_Helper(X=input, W=W, R=R, B=B, direction='reverse')
        output = gru.run()
        print_results(output)

    @staticmethod
    def Legacy_TestGRUOpBidirectionalBasic():

        print(GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpBidirectionalBasic.__name__)

        input = np.array([[[-0.455351, -0.276391]],
                          [[-0.185934, -0.269585]]]).astype(np.float32)

        W, R, B = ONNXRuntimeTestContext.BidirectionalWeights()
        gru = GRU_Helper(X=input, W=W, R=R, B=B, direction='bidirectional')
        output = gru.run()
        print_results(output)

GRU_ONNXRuntimeUnitTests.ForwardDefaultActivationsSimpleWeightsNoBiasTwoRows()
GRU_ONNXRuntimeUnitTests.ReverseDefaultActivationsSimpleWeightsNoBiasTwoRows()
GRU_ONNXRuntimeUnitTests.BidirectionalDefaultActivationsSimpleWeightsNoBiasTwoRows()

GRU_ONNXRuntimeUnitTests.ForwardDefaultActivationsSimpleWeightsWithBiasBatchParallel()
GRU_ONNXRuntimeUnitTests.ForwardDefaultActivationsSimpleWeightsWithBiasBatchParallelLinearBeforeReset()
GRU_ONNXRuntimeUnitTests.ReverseDefaultActivationsSimpleWeightsWithBiasBatchParallelLinearBeforeReset()
GRU_ONNXRuntimeUnitTests.ForwardDefaultActivationsSimpleWeightsWithBiasLinearBeforeReset()
GRU_ONNXRuntimeUnitTests.ReverseDefaultActivationsSimpleWeightsWithBiasLinearBeforeReset()

GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpForwardBasic()
GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpBackwardBasic()
GRU_ONNXRuntimeUnitTests.Legacy_TestGRUOpBidirectionalBasic()
