import sys
import os
from unit_test import UnitTest
import pytest
import torch
import torch_mlu
import btunittests

import math

def alibi_slopes(batch,
                 num_heads,
                 true_seq_lens,
                 train_seq_len,
                 use_dynamic):
    scale = 1.0
    if use_dynamic:
        # dynamic ntk factor according to actual sequence length
        a0 = 1.0
        # train_seq_len = 2048
        dynamic_seq_len = true_seq_lens  # [batch, 1]
        a = a0 * dynamic_seq_len / train_seq_len  # [batch, 1]
        a = a.masked_fill(a < 1.0, 1.0)  # dynamic step 1: dynamic ntk scaling factor
        scale = a ** (1.0 / (num_heads-1))  # dynamic step 2: coefficient b, for computation convenience
    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
    base = torch.tensor(
        2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), dtype=torch.float32
    )
    if use_dynamic:
        base = base / scale  # dynamic step 3: divide b to alibi base
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)
    if use_dynamic:
        slopes = slopes * scale  # dynamic step 4: fix alibi bias m_h by multiplying b

    if closest_power_of_2 != num_heads:  # todo: fix ntk when num_heads is not power of 2
        extra_base = torch.tensor(
            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), dtype=torch.float32
        )
        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
        extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, dtype=torch.int32)
        extra_slopes = torch.pow(extra_base, extra_powers)
        if use_dynamic:
            extra_slopes = extra_slopes.unsqueeze(0)
            extra_slopes = extra_slopes.repeat(batch, 1)
        slopes = torch.cat([slopes, extra_slopes], dim=-1)

    if not use_dynamic:
        slopes = slopes.unsqueeze(0)
        slopes = slopes.repeat(batch, 1)

    return slopes

class TestAlibiSlopeKernel(UnitTest):
    @pytest.mark.parametrize("batch, tp_head_num, tp_num, use_dynamic, max_sequence_length", [
        (1, 11, 1, False, 1024),
        (1, 11, 2, False, 1024),
        (8, 16, 1, False, 4096),
        (8, 16, 2, False, 4096),
        (1, 11, 1, True, 1024),
        (1, 11, 2, True, 1024),
        (8, 16, 1, True, 4096),
        (8, 16, 2, True, 4096)])
    def test_alibi_slope(self, batch, tp_head_num, tp_num, use_dynamic, max_sequence_length):
        torch.manual_seed(0)
        true_seq_lens = torch.randint(0, 2 * max_sequence_length, (batch, 1), dtype=torch.int32)
        output_base = alibi_slopes(batch, tp_head_num * tp_num, true_seq_lens, max_sequence_length, use_dynamic)
        output_mlu = btunittests.alibi_slope_test(true_seq_lens.mlu(), batch, tp_head_num, tp_num, use_dynamic, max_sequence_length)
        diff1 = super().diff1(output_mlu, output_base)
        diff2 = super().diff2(output_mlu, output_base)
        assert diff1 <= 1e-5 and diff2 <= 1e-5

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)
