import argparse
import numpy as np
import torch
import torch.nn.functional as F
import pytest
import json
import random
import time
import copy
from torch import nn
from modelingpreln import BertEncoder as BertEncoderPreln
from modeling import BertEncoder as BertEncoderPostln
from modeling import BertLayerNorm, BertConfig
from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig
import deepspeed

import sys

if not deepspeed.ops.__installed_ops__['transformer']:
    pytest.skip("transformer kernels are not installed", allow_module_level=True)


def check_equal(first, second, atol=1e-2, verbose=False):
    if verbose:
        print()
    for i, (x, y) in enumerate(zip(first, second)):
        x = x[0].cpu().detach().numpy()
        y = y[0].cpu().detach().numpy()
        if verbose:
            print("x = {}".format(x.flatten()))
            print("y = {}".format(y.flatten()))
            print('-' * 80)
        np.testing.assert_allclose(x, y, err_msg="Index: {}".format(i), atol=atol)


def zero_grad(variables):
    for variable in variables:
        variable.grad.zero_()


device = torch.device("cuda")
kwargs_fp32 = {'dtype': torch.float, 'device': device, 'requires_grad': True}
kwargs_fp16 = {'dtype': torch.half, 'device': device, 'requires_grad': True}


class DSEncoder(nn.Module):
    def __init__(self, config, weights, biases):
        super(DSEncoder, self).__init__()
        self.FinalLayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.layer = nn.ModuleList([
            copy.deepcopy(DeepSpeedTransformerLayer(i,
                                                    config,
                                                    weights,
                                                    biases))
            for i in range(config.num_hidden_layers)
        ])
        self.grads = []
        self.pre_or_post = config.pre_layer_norm

    def forward(self,
                hidden_states,
                attention_mask,
                output_all_encoded_layers=True,
                checkpoint_activations=False):
        all_encoder_layers = []

        def custom(start, end):
            def custom_forward(*inputs):
                layers = self.layer[start:end]
                x_ = inputs[0]
                for layer in layers:
                    x_ = layer(x_, inputs[1])
                return x_

            return custom_forward

        if checkpoint_activations:
            l = 0
            num_layers = len(self.layer)
            chunk_length = math.ceil(math.sqrt(num_layers))
            while l < num_layers:
                hidden_states = checkpoint.checkpoint(custom(l,
                                                             l + chunk_length),
                                                      hidden_states,
                                                      attention_mask * 1)
                l += chunk_length
            # decoder layers
        else:
            for i, layer_module in enumerate(self.layer):
                hidden_states = layer_module(hidden_states, attention_mask)
                hidden_states.register_hook(
                    lambda x,
                    i=i,
                    self=self: self.grads.append([x,
                                                  "hidden_state"]))

                if output_all_encoded_layers:
                    all_encoder_layers.append(hidden_states)

        if not output_all_encoded_layers or checkpoint_activations:
            if (self.pre_or_post):
                hidden_states = self.FinalLayerNorm(hidden_states)
            all_encoder_layers.append(hidden_states)
        return all_encoder_layers

    def get_grads(self):
        return self.grads


def create_models(ds_config):
    bert_config = BertConfig(vocab_size_or_config_json_file=119547,
                             hidden_size=ds_config.hidden_size,
                             num_hidden_layers=ds_config.num_hidden_layers,
                             num_attention_heads=ds_config.heads,
                             batch_size=ds_config.batch_size,
                             intermediate_size=ds_config.intermediate_size,
                             hidden_act="gelu",
                             hidden_dropout_prob=ds_config.hidden_dropout_ratio,
                             attention_probs_dropout_prob=ds_config.attn_dropout_ratio,
                             max_position_embeddings=ds_config.max_seq_length,
                             type_vocab_size=2,
                             initializer_range=ds_config.initializer_range,
                             fp16=ds_config.fp16)

    weights = []
    biases = []

    for i in range(4):
        weights.append(
            nn.Parameter(torch.Tensor(ds_config.hidden_size,
                                      ds_config.hidden_size)))
        weights[i].data.normal_(mean=0.0, std=ds_config.initializer_range)

    weights.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
    weights[4].data.fill_(1.0)
    weights.append(
        nn.Parameter(torch.Tensor(ds_config.intermediate_size,
                                  ds_config.hidden_size)))
    weights[5].data.normal_(mean=0.0, std=ds_config.initializer_range)
    weights.append(
        nn.Parameter(torch.Tensor(ds_config.hidden_size,
                                  ds_config.intermediate_size)))
    weights[6].data.normal_(mean=0.0, std=ds_config.initializer_range)
    weights.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
    weights[7].data.fill_(1.0)

    biases.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
    biases[0].data.zero_()
    for i in range(4):
        biases.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
        biases[i + 1].data.zero_()
    biases.append(nn.Parameter(torch.Tensor(ds_config.intermediate_size)))
    biases[5].data.zero_()
    biases.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
    biases[6].data.zero_()
    biases.append(nn.Parameter(torch.Tensor(ds_config.hidden_size)))
    biases[7].data.zero_()

    if (ds_config.pre_layer_norm):
        bert_encoder = BertEncoderPreln(bert_config, weights, biases)
    else:
        bert_encoder = BertEncoderPostln(bert_config, weights, biases)
    ds_encoder = DSEncoder(ds_config, weights, biases)

    if ds_config.fp16:
        bert_encoder.half()
        ds_encoder.half()

    bert_encoder.cuda()
    ds_encoder.cuda()

    return bert_encoder, ds_encoder


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def run_forward(ds_config, seq_len, atol=1e-2, verbose=False, test_bsz=None):
    set_seed(123)
    bert_encoder, ds_encoder = create_models(ds_config)

    bsz = ds_config.batch_size if test_bsz is None else test_bsz

    # prepare test data
    kwargs = kwargs_fp16 if ds_config.fp16 else kwargs_fp32
    hidden_states = torch.randn(bsz,
                                seq_len, #ds_config.max_seq_length,
                                ds_config.hidden_size,
                                **kwargs)
    input_mask = torch.randn(bsz, 1, 1,
                             seq_len, #ds_config.max_seq_length,
                             **kwargs)

    # run baseline
    base_results = bert_encoder(hidden_states,
                                input_mask,
                                output_all_encoded_layers=False,
                                checkpoint_activations=False)

    # run ds
    ds_results = ds_encoder(hidden_states,
                            input_mask,
                            output_all_encoded_layers=False,
                            checkpoint_activations=False)

    # check grads
    check_equal(base_results, ds_results, atol=atol, verbose=verbose)


# FP16 test cases can only run on the devices support FP16.
@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
                         [
                             (64,1024,128,16,3,True,False),
                             (64,1024,128,16,3,True,True),
                             (8,1024,384,16,3,True,False),
                             (8,1024,384,16,3,True,True),
                             (8,1024,384,16,3,True,True),
                             (8,1024,120,16,3,True,False),
                             (8,1024,120,16,3,True,True),
                             (8,1024,512,16,3,True,False),
                             (8,1024,512,16,3,True,True),
                             (64,1024,56,16,3,False,False),
                             (64,1024,56,16,3,False,True),
                             (64,1024,24,16,3,False,False),
                             (64,1024,24,16,3,False,True),
                             (8,1024,384,16,3,False,False),
                             (8,1024,384,16,3,False,True),
                             (8,1024,512,16,3,False,False),
                             (8,1024,512,16,3,False,True),
                             (8,1536,128,24,3,False,False),
                             (8,1536,128,24,3,False,True),
                             (8,2048,128,32,3,False,False),
                             (8,2048,128,32,3,False,True),
                             (8,2560,128,40,3,False,False),
                             (8,2560,128,40,3,False,True),
                         ]) # yapf: disable
def test_forward(batch_size,
                 hidden_size,
                 seq_len,
                 heads,
                 num_layers,
                 is_preln,
                 use_fp16):
    # Only run fp16 test cases on devices with 7+ capability.
    major, _ = torch.cuda.get_device_capability()
    if major < 7 and use_fp16 is True:
        return

    ds_config = DeepSpeedTransformerConfig()
    ds_config.layer_id = None
    ds_config.batch_size = batch_size
    ds_config.hidden_size = hidden_size
    ds_config.max_seq_length = 128  #seq_len
    ds_config.intermediate_size = 4 * hidden_size
    ds_config.heads = heads
    ds_config.attn_dropout_ratio = 0.0
    ds_config.hidden_dropout_ratio = 0.0
    ds_config.num_hidden_layers = num_layers
    ds_config.pre_layer_norm = is_preln
    ds_config.initializer_range = 0.02
    ds_config.fp16 = use_fp16

    run_forward(ds_config, seq_len, atol=2e-2)


@pytest.mark.parametrize('batch_size, small_bsz, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
                         [
                             (8,3,1024,512,16,3,True,False),
                             (8,7,1024,512,16,3,True,True),
                             (8,3,1024,512,16,3,False,False),
                             (8,7,1024,512,16,3,False,True),
                         ]) # yapf: disable
def test_forward_with_small_bsz(batch_size,
                                small_bsz,
                                hidden_size,
                                seq_len,
                                heads,
                                num_layers,
                                is_preln,
                                use_fp16):
    # Only run fp16 test cases on devices with 7+ capability.
    major, _ = torch.cuda.get_device_capability()
    if major < 7 and use_fp16 is True:
        return

    ds_config = DeepSpeedTransformerConfig()
    ds_config.layer_id = None
    ds_config.batch_size = batch_size
    ds_config.hidden_size = hidden_size
    ds_config.intermediate_size = 4 * hidden_size
    ds_config.max_seq_length = seq_len
    ds_config.heads = heads
    ds_config.attn_dropout_ratio = 0.0
    ds_config.hidden_dropout_ratio = 0.0
    ds_config.num_hidden_layers = num_layers
    ds_config.pre_layer_norm = is_preln
    ds_config.initializer_range = 0.02
    ds_config.fp16 = use_fp16

    run_forward(ds_config, seq_len, atol=2e-2, test_bsz=small_bsz)

@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
                         [
                             (64,1024,128,16,3,True,False),
                             (64,1024,128,16,3,True,True),
                             (64,1024,128,16,3,False,False),
                             (64,1024,128,16,3,False,True),
                         ]) # yapf: disable
def test_forward_stochastic(batch_size,
                            hidden_size,
                            seq_len,
                            heads,
                            num_layers,
                            is_preln,
                            use_fp16):
    # Only run fp16 test cases on devices with 7+ capability.
    major, _ = torch.cuda.get_device_capability()
    if major < 7 and use_fp16 is True:
        return

    ds_config = DeepSpeedTransformerConfig()
    ds_config.layer_id = None
    ds_config.batch_size = batch_size
    ds_config.hidden_size = hidden_size
    ds_config.intermediate_size = 4 * hidden_size
    ds_config.max_seq_length = seq_len
    ds_config.heads = heads
    ds_config.attn_dropout_ratio = 0.0
    ds_config.hidden_dropout_ratio = 0.0
    ds_config.num_hidden_layers = num_layers
    ds_config.pre_layer_norm = is_preln
    ds_config.initializer_range = 0.02
    ds_config.fp16 = use_fp16
    ds_config.stochastic_mode = True

    run_forward(ds_config, seq_len, atol=7e-2)