import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from unit_test import UnitTest
import pytest
import torch
import torch_mlu
import btunittests
import random
import math

def create_cos_sin_table(rotary_base, rotary_seq_len_, rotary_stride_, rotary_scaling, rotary_dim, interleaved, datatype):
    torch.manual_seed(0)
    random.seed(0)
    inv_freq=1.0 / (rotary_base ** (torch.arange(0, rotary_dim, 2).mlu() / rotary_dim))
    t = torch.arange(rotary_seq_len_).mlu()
    t = t / rotary_scaling
    freqs = torch.outer(t, inv_freq).mlu()
    cos_table = torch.cos(freqs).mlu()
    sin_table = torch.sin(freqs).mlu()
    if interleaved == 0:
        cos_table = torch.cat((cos_table,cos_table), dim = 1).mlu()
        sin_table = torch.cat((sin_table,sin_table), dim = 1).mlu()
        emb = torch.cat((cos_table, sin_table), dim=-1).mlu()
    else:
        cos_tmp = torch.cat((cos_table[:,0].unsqueeze(1),cos_table[:,0].unsqueeze(1)), dim = 1).mlu()
        emb = cos_tmp
        for i in range(1, cos_table.size(1)):
            cos_tmp = torch.cat((cos_table[:,i].unsqueeze(1),cos_table[:,i].unsqueeze(1)), dim = 1).mlu()
            emb = torch.cat((emb,cos_tmp), dim = 1).mlu()
        for i in range(0, cos_table.size(1)):
            sin_tmp = torch.cat((sin_table[:,i].unsqueeze(1),sin_table[:,i].unsqueeze(1)), dim = 1).mlu()
            emb = torch.cat((emb,sin_tmp), dim = 1).mlu()
    return emb

def get_ntk_alpha(true_seq_len, max_position_embeddings, rotary_base, rotary_dim):
    context_value = math.log(true_seq_len / max_position_embeddings, 2) + 1
    ntk_alpha = 2 ** math.ceil(context_value) - 1
    ntk_alpha = max(ntk_alpha, 1)
    rotary_b = rotary_base * (ntk_alpha ** (rotary_dim / (rotary_dim - 2)))
    return ntk_alpha, rotary_b

class TestCreateCosSinTableKernel(UnitTest):
    @pytest.mark.parametrize("batch, max_position_embeddings, batch_stride, rotary_seq_len, rotary_stride, rotary_base, rotary_scaling,\
    rotary_scaling_type, rotary_dim , interleaved, datatype",
    [(8, 512, 0, 267, 128, 10000, 1, 1, 64, 0, torch.float32),
     (8, 512, 0, 267, 128, 10000, 1, 1, 64, 1, torch.float32),
     (4, 256, 11456, 179, 64, 10000, 1, 2, 32, 0, torch.float32),
     (4, 256, 11456, 179, 64, 10000, 1, 2, 32, 1, torch.float32),
     (4, 256, 34368, 179, 192, 10000, 1, 2, 96, 0, torch.float32),
     (4, 256, 34368, 179, 192, 10000, 1, 2, 96, 1, torch.float32),
     (4, 256, 45824, 179, 256, 10000, 1, 2, 128, 0, torch.float32),
     (4, 256, 45824, 179, 256, 10000, 1, 2, 128, 1, torch.float32),
     (8, 512, 0, 267, 128, 10000, 1, 1, 64, 0, torch.bfloat16),
     (8, 512, 0, 267, 128, 10000, 1, 1, 64, 1, torch.bfloat16),
     (4, 256, 11456, 179, 64, 10000, 1, 2, 32, 0, torch.float16),
     (4, 256, 11456, 179, 64, 10000, 1, 2, 32, 1, torch.float16),
     (4, 256, 34368, 179, 192, 10000, 1, 2, 96, 0, torch.float16),
     (4, 256, 34368, 179, 192, 10000, 1, 2, 96, 1, torch.float16),
     (4, 256, 45824, 179, 256, 10000, 1, 2, 128, 0, torch.float16),
     (4, 256, 45824, 179, 256, 10000, 1, 2, 128, 1, torch.float16)])
    def test_all(self, batch, max_position_embeddings, batch_stride, rotary_seq_len, rotary_stride, rotary_base, rotary_scaling,
                 rotary_scaling_type, rotary_dim, interleaved ,datatype):
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name and datatype == torch.bfloat16:
            datatype = torch.half
            print("BFLOAT16 is not support on {}, use half instead".format(mlu_name))
        torch.manual_seed(0)
        random.seed(0)
        seq_lens = torch.randint(size = (batch,), low = 0, high = 512,
                                     dtype = torch.int32).mlu()
        ntk_alpha_list = []
        rotary_emb_alpha_cached = torch.rand((batch), dtype=torch.float32).mlu()
        ref_rotary_emb_alpha_cached = rotary_emb_alpha_cached.clone()
        output_mlu = btunittests.create_cos_sin_table_test(ref_rotary_emb_alpha_cached, seq_lens, max_position_embeddings, batch_stride,
                                                           rotary_seq_len, rotary_dim, rotary_stride, rotary_scaling,
                                                           rotary_scaling_type, rotary_base, interleaved)
        if rotary_scaling_type !=  2:
            output_base = create_cos_sin_table(rotary_base, rotary_seq_len, rotary_stride, rotary_scaling, rotary_dim, interleaved, datatype)
        else:
            max_seq_len, max_index = torch.max(seq_lens, dim = 0)
            if max_seq_len > max_position_embeddings:
                for i in range(batch):
                    ntk_alpha, rb = get_ntk_alpha(seq_lens[i], max_position_embeddings, rotary_base, rotary_dim)
                    rotary_emb_alpha_cached[i] = ntk_alpha
                    cos_sin_table_tmp = create_cos_sin_table(rb, rotary_seq_len, rotary_stride, rotary_scaling, rotary_dim, interleaved, datatype).unsqueeze(0)
                    cos_sin_table = torch.cat((cos_sin_table,cos_sin_table_tmp),dim = 0).mlu() if i != 0 else cos_sin_table_tmp
            else:
                ntk_alpha, rb = get_ntk_alpha(max_seq_len, max_position_embeddings, rotary_base, rotary_dim)
                for i in range(batch):
                    rotary_emb_alpha_cached[i] = ntk_alpha
                    cos_sin_table_tmp = create_cos_sin_table(rb, rotary_seq_len, rotary_stride, rotary_scaling, rotary_dim, interleaved, datatype).unsqueeze(0)
                    cos_sin_table = torch.cat((cos_sin_table,cos_sin_table_tmp),dim = 0).mlu() if i != 0 else cos_sin_table_tmp
            output_base = cos_sin_table
            output_base_1 = rotary_emb_alpha_cached
        diff1 = super().diff1(output_mlu, output_base)
        diff2 = super().diff2(output_mlu, output_base)
        assert diff1 < 0.003 and diff2 < 0.003
        if rotary_scaling_type == 2:
            diff1 = super().diff1(ref_rotary_emb_alpha_cached, output_base_1)
            diff2 = super().diff2(ref_rotary_emb_alpha_cached, output_base_1)
            assert diff1 < 0.003 and diff2 < 0.003
            assert diff1 < 0.003 and diff2 < 0.003

if __name__ == '__main__':
    exit_code = pytest.main(['-vs', os.path.abspath(__file__), '--junitxml=' + os.path.splitext(os.path.basename(__file__))[0] + '_report.xml'])
    exit(exit_code)
