import triton

import os
import sys
import pdb
# sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
import math
import argparse

import torch
import torch.nn.functional as F
import pytest
import numpy
from unittest import mock
from megatron.model.transformer import ParallelAttention, BSMMParallelAttention
from megatron import get_args
from megatron.initialize import initialize_megatron

from megatron.model.enums import AttnMaskType, AttnType
from megatron.model import LayerNorm
from megatron.model.utils import init_method_normal

def assert_almost_equal(x, y, decimal=2, err_msg=''):
    import numpy.testing as npt
    if isinstance(x, torch.Tensor):
        x = x.cpu().detach().numpy()
    if isinstance(y, torch.Tensor):
        y = y.cpu().detach().numpy()
    npt.assert_array_almost_equal(x, y, err_msg=err_msg, decimal=decimal)

def test_attn():
    initialize_megatron(extra_args_provider=None, args_defaults={'tokenizer_type': 'GPT2BPETokenizer','no_load_rng': True, 'no_load_optim': True})
    args = get_args()

    init_method = init_method_normal(args.init_method_std)
    output_layer_init_method = init_method_normal(args.init_method_std)
    layer_number = 1
    #######################
    # self attention test #
    #######################
    self_attn = ParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.self_attn,
            attn_mask_type=AttnMaskType.padding
            )
    torch.save(self_attn.state_dict(), "./self_attn.pt")
    # print(self_attn.state_dict())
    # pdb.set_trace()

    self_attn_bsmm = BSMMParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.self_attn,
            attn_mask_type=AttnMaskType.padding)
    self_attn_bsmm.load_state_dict(torch.load("./self_attn.pt"))
    
    # hidden_states: [sq, b, h]... 그래서 이게 key야 query야 value야??
    hidden_states = torch.randn(args.seq_length, args.micro_batch_size, args.hidden_size).float().cuda().contiguous()
    attention_mask = torch.tril(torch.ones(args.micro_batch_size, args.seq_length, args.seq_length)).view(args.micro_batch_size, 1, args.seq_length, args.seq_length).cuda().contiguous()
    attention_mask = (attention_mask < 0.5).cuda().contiguous()
    attention_mask_bsmm = torch.tril(torch.ones(args.seq_length, args.seq_length)).cuda()

    # save query_layer, 
    encoder_output, encoder_bias = self_attn(hidden_states, attention_mask, layer_past=None, get_key_value=False)
    encoder_output_bsmm, encoder_bias_bsmm = self_attn_bsmm(hidden_states, attention_mask_bsmm, layer_past=None, get_key_value=False)

    assert_almost_equal(encoder_output, encoder_output_bsmm)
    assert_almost_equal(encoder_bias, encoder_bias_bsmm)
    # ########################
    # # cross attention test #
    ########################
    cross_attn = ParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.cross_attn,
            attn_mask_type=AttnMaskType.causal)
    torch.save(cross_attn.state_dict(), "./cross_attn.pt")
    
    cross_attn_bsmm = BSMMParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.cross_attn,
            attn_mask_type=AttnMaskType.causal)
    cross_attn_bsmm.load_state_dict(torch.load("./cross_attn.pt"))
            
    hidden_states = torch.randn(args.seq_length, args.micro_batch_size, args.hidden_size).float().cuda().contiguous()
    attention_mask = torch.tril(torch.ones(args.micro_batch_size, args.seq_length, args.seq_length)).view(args.micro_batch_size, 1, args.seq_length, args.seq_length).cuda().contiguous()
    attention_mask = (attention_mask < 0.5).cuda().contiguous()
    attention_mask_bsmm = torch.tril(torch.ones(args.seq_length, args.seq_length)).cuda()

    decoder_output, decoder_bias = cross_attn(hidden_states, attention_mask, layer_past=None, get_key_value=False, encoder_output=encoder_output)

    decoder_output_bsmm, decoder_bias_bsmm = cross_attn_bsmm(hidden_states,
            attention_mask_bsmm,
            layer_past=None,
            get_key_value=False,
            encoder_output=encoder_output_bsmm)

    assert_almost_equal(decoder_output, decoder_output_bsmm)
    assert_almost_equal(decoder_bias, decoder_bias_bsmm)


if __name__ == "__main__":
        test_attn()