import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from torch.nn.parameter import Parameter
from torch.nn import functional as F
import torch.nn as nn

class SelfAttn(torch.nn.Module):
    def __init__(
        self,
        qkv_weights,
        qkv_biass,
        o_weight,
        o_bias,
        norm_weight,
        norm_bias,
        input_size,
        head_size,
        query_factor,
        eps = 1e-5
    ) -> None:
        super().__init__()
        assert len(qkv_weights) == 3 and len(qkv_biass) == 3, 'length of weights and biass must be 3'
        self.layernorm = torch.nn.LayerNorm(input_size)
        self.layernorm.eps = eps
        self.layernorm.weight = nn.Parameter(norm_weight)
        self.layernorm.bias = nn.Parameter(norm_bias)
        self.weights = qkv_weights
        self.biass = qkv_biass
        self.o_weight  = o_weight
        self.o_bias = o_bias
        self.head_size = head_size
        self.head_num = self.weights[0].size(0) // head_size
        self.query_factor = query_factor

    def forward(self, input: torch.Tensor):
        n = input.size(0)
        t = input.size(1)
        normed_input = self.layernorm(input)
        q = F.linear(normed_input, nn.Parameter(self.weights[0]), nn.Parameter(self.biass[0])).view(n, t, self.head_num, self.head_size)
        k = F.linear(normed_input, nn.Parameter(self.weights[1]), nn.Parameter(self.biass[1])).view(n, t, self.head_num, self.head_size)
        v = F.linear(normed_input, nn.Parameter(self.weights[2]), nn.Parameter(self.biass[2])).view(n, t, self.head_num, self.head_size)
        qk = torch.einsum('bthd,bshd->bhts', q, k) * self.query_factor
        attn = torch.softmax(qk, dim=-1, dtype=v.dtype)
        qkv = torch.einsum('bhts,bshd->bthd', attn, v).reshape(n, t, -1)
        output = F.linear(qkv, nn.Parameter(self.o_weight), nn.Parameter(self.o_bias)) + input
        return output

class BTSelfAttn(torch.nn.Module):
    def __init__(
        self,
        qkv_weights,
        qkv_biass,
        o_weight,
        o_bias,
        norm_weight,
        norm_bias,
        head_size,
        query_factor,
        eps = 1e-5
    ) -> None:
        super().__init__()
        self.weights = qkv_weights
        self.biass = qkv_biass
        self.o_weight = o_weight
        self.o_bias = o_bias
        self.norm_weight = norm_weight
        self.norm_bias = norm_bias
        self.head_size = head_size
        self.query_factor = query_factor
        self.eps = eps

    def forward(self, input: torch.Tensor):
        n, t = input.size(0), input.size(1)
        q, k, v =  ops.fused_norm_attention_project(input,
                             self.weights[0],
                             self.biass[0],
                             self.weights[1],
                             self.biass[1],
                             self.weights[2],
                             self.biass[2],
                             self.norm_weight,
                             self.norm_bias,
                             self.eps,
                             "nhtc",
                             self.head_size,
                             False)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        attn_out = ops.flash_attention(q, k, v, None, None, None, None, None, t, t,
                                          self.query_factor, False, -1, -1, q.dtype).flatten(-2, -1)
        output = ops.attention_project(attn_out, self.o_weight, self.o_bias, input)
        return output

dtype_list = [torch.half, torch.float]
if torch_mlu.mlu.is_bf16_supported():
    dtype_list.append(torch.bfloat16)
class TestSelfAttn(BtTestCase):
    def op_impl_base(self, *args):
        return super().op_impl_base(*args)

    def test_self_attn(self):
        N, T, input_size, hidden_size, head_size, eps, query_factor = 5, 2048, 512, 768, 64, 1e-5, 0.125
        for dtype in dtype_list:
            print("N: {}, T: {}, input_size: {}, hidden_size: {}, testing...".format(
                    N, T, input_size, hidden_size, dtype), flush=True)
            input = torch.randn(N, T, input_size, dtype=dtype, device="mlu")
            weight = torch.randn(hidden_size * 3, input_size, dtype=dtype, device="mlu") / 10
            bias = torch.randn(hidden_size * 3, dtype=dtype, device="mlu")
            norm_weight = torch.randn(input_size, dtype=dtype, device="mlu")
            norm_bias = torch.randn(input_size, dtype=dtype, device="mlu")
            weights = torch.chunk(weight, 3)
            biass = torch.chunk(bias, 3)
            o_weight = torch.randn(input_size, hidden_size, dtype=dtype, device="mlu") / 10
            o_bias = torch.randn(input_size, dtype=dtype, device="mlu")

            torch_self_attn = SelfAttn(weights, biass, o_weight, o_bias, norm_weight, norm_bias,
                                    input_size, head_size, query_factor, eps)
            tmo_self_attn = BTSelfAttn(weights, biass, o_weight, o_bias, norm_weight, norm_bias,
                                    head_size, query_factor, eps)

            # test self_attn
            print("test self_attn...")
            torch_out = torch_self_attn(input)
            tmo_out = tmo_self_attn(input)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                        0.011, use_MSE=True, use_RAE=True)

    def test_sd_attention(self):
        N, Tq, Tk, hq, hk, head_size, query_factor = 2, 4096, 77, 8, 8, 40, 0.125
        for dtype in dtype_list:
            q = torch.randn(N, Tq, hq, head_size, dtype=dtype, device="mlu")
            k = torch.randn(N, Tk, hk, head_size, dtype=dtype, device="mlu")
            v = torch.randn(N, Tk, hk, head_size, dtype=dtype, device="mlu")

            qk = torch.einsum('bthd,bshd->bhts', q, k) * query_factor
            attn = torch.softmax(qk, dim=-1, dtype=v.dtype)
            torch_out = torch.einsum('bhts,bshd->bthd', attn, v).reshape(N, Tq, -1)

            qt = q.transpose(1, 2).contiguous()
            kt = k.transpose(1, 2).contiguous()
            vt = v.transpose(1, 2).contiguous()
            tmo_out = ops.flash_attention(qt.transpose(1, 2),
                                        kt.transpose(1, 2),
                                        vt.transpose(1, 2),
                                        None, None, None, None, None,
                                        Tq, Tk, query_factor, False).flatten(-2, -1)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                    0.004, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        return super().test_inductor()

if __name__ == '__main__':
    exit(run_unittest(TestSelfAttn))
