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

active = torch.nn.functional.silu

def torch_ffn(input, w1, scale1, bias1, w2, scale2, bias2, is_gated):
    tmp_input = input.flatten(0, -2).to(torch.float)
    inner_size = w1.size(0) // (1 + is_gated)
    imm = weight_only_quant_matmul(tmp_input, w1, scale1, bias1)
    acted = active(imm[:, :inner_size])
    acted = acted * imm[:, inner_size:] if is_gated else acted
    out = weight_only_quant_matmul(acted, w2, scale2, bias2)
    return out.reshape(input.shape)

def torch_smooth_quant_ffn(input, w1, scale1, bias1, w2, scale2, bias2, input_smooth, act_smooth, is_gated):
    inner_size = w1.size(0) // (1 + is_gated)
    tmp_input = input.flatten(0, -2).to(torch.float)
    quant_input, input_scale = QuantByRow(tmp_input.flatten(0, -2) * input_smooth, 8)
    imm = smooth_quant_matmul(quant_input, input_scale, w1, scale1, input.dtype, bias1)
    acted = active(imm[:, :inner_size])
    acted = acted * imm[:, inner_size:] if is_gated else acted
    quant_acted, acted_scale = QuantByRow(acted.flatten(0, -2) * act_smooth, 8)
    out = smooth_quant_matmul(quant_acted, acted_scale, w2, scale2, input.dtype, bias2)
    return out.reshape(input.shape)

def tmo_weight_only_quant_ffn(input, w1, scale1, bias1, w2, scale2, bias2, quant_bit):
    tmp_input = input.flatten(0, -2)
    imm = ops.weight_only_quant_matmul(tmp_input, w1, scale1, None, bias1, None, 'silu', quant_bit, True)
    out = ops.weight_only_quant_matmul(imm, w2, scale2, None, bias2, None, "none", quant_bit)
    return out.reshape(input.shape)

def tmo_weight_only_group_quant_ffn(input, w1, scale1, bias1, w2, scale2, bias2, quant_bit):
    tmp_input = input.flatten(0, -2)
    imm = ops.weight_only_quant_matmul(tmp_input, w1, scale1, None, bias1, None, 'none', quant_bit)
    acted = active(imm)
    out = ops.weight_only_quant_matmul(acted, w2, scale2, None, bias2, None, "none", quant_bit)
    return out.reshape(input.shape)

def tmo_weight_only_quant_gated_ffn(input, w1, scale1, bias1, w2, scale2, bias2, quant_bit):
    tmp_input = input.flatten(0, -2)
    imm = ops.weight_only_quant_matmul(tmp_input, w1, scale1, None, bias1, None, 'none', quant_bit)
    acted = ops.active(imm, 'silu', True)
    out = ops.weight_only_quant_matmul(acted, w2, scale2, None, bias2, None, "none", quant_bit)
    return out.reshape(input.shape)

def tmo_pertoken_smooth_quant_gated_ffn(input, w1, scale1, bias13, w2, scale2, bias2, input_smooth, act_smooth, dtype):
    tmp_input = input.flatten(0, -2)
    quant_input, input_scale = ops.per_token_smooth_quantize(tmp_input, input_smooth, None)
    imm = ops.smooth_quant_matmul(quant_input, input_scale, w1, scale1, dtype, bias13)
    acted = ops.active(imm, 'silu', True)
    quant_acted, acted_scale = ops.per_token_smooth_quantize(acted, act_smooth, None)
    out = ops.smooth_quant_matmul(quant_acted, acted_scale, w2, scale2, dtype, bias2)
    return out.reshape(input.shape)

def init_tensors(batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype=torch.half, group_num=1):
    sigma = 0.1
    eps = 0.1       # Avoid the occurrence of nan
    torch.manual_seed(1)
    input_smooth = torch.randn(hidden_size, dtype=torch.float, device="mlu").abs() + eps
    act_smooth = torch.randn(inner_size, dtype=torch.float, device="mlu").abs() + eps
    w1 = torch.randn((1 + is_gated) * inner_size, hidden_size, dtype=dtype, device="mlu") * sigma
    w1 = w1 / input_smooth if with_smooth else w1
    bias1 = torch.randn((1 + is_gated) * inner_size, dtype=dtype, device="mlu") * sigma
    w2 = torch.randn(hidden_size, inner_size, dtype=dtype, device="mlu") * sigma
    w2 = w2 / act_smooth if with_smooth else w2
    bias2 = torch.randn(hidden_size, dtype=dtype, device="mlu") * sigma
    input = torch.randn(batch, seq, hidden_size, dtype=dtype, device="mlu")
    quant_w1, scale1 = QuantByRow(w1, quant_bit, group_num)
    quant_w2, scale2 = QuantByRow(w2, quant_bit, group_num)
    return input, quant_w1, scale1, quant_w2, scale2, bias1, bias2, input_smooth, act_smooth

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

    def test_weight_only_quant_ffn(self):
        for dtype in dtype_list:
            print("test_weight_only_quant_ffn...")
            batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth  = 3, 5, 512, 768, 8, False, False
            input, quant_w1, scale1, quant_w2, scale2, bias1, bias2, _, _ = init_tensors(batch, seq,
                    hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype=dtype)
            torch_out = torch_ffn(input, quant_w1, scale1, bias1, quant_w2, scale2, bias2, False)
            tmo_out = tmo_weight_only_quant_ffn(input, quant_w1, scale1, bias1, quant_w2, scale2, bias2, quant_bit)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                    0.005, use_MSE=True, use_RAE=True)

    def test_weight_only_quant_gated_ffn(self):
        for dtype in dtype_list:
            print("test_weight_only_quant_gated_ffn...")
            batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth  = 3, 5, 512, 768, 4, True, False
            input, quant_w1, scale1, quant_w2, scale2, bias1, bias2, _, _ = init_tensors(batch, seq,
                    hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype=dtype)
            quant_w1_int4 = PairlyPackInt8(quant_w1)
            quant_w2_int4 = PairlyPackInt8(quant_w2)
            torch_out = torch_ffn(input, quant_w1, scale1, bias1, quant_w2, scale2, bias2, True)
            tmo_out = tmo_weight_only_quant_gated_ffn(input, quant_w1_int4, scale1, bias1, quant_w2_int4, scale2, bias2, quant_bit)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                    0.005, use_MSE=True, use_RAE=True)

    def test_weight_only_group_quant_ffn(self):
        for dtype in dtype_list:
            print("test_weight_only_group_quant_ffn...")
            batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth  = 3, 5, 512, 512, 8, False, False
            input, quant_w1, scale1, quant_w2, scale2, bias1, bias2, _, _ = init_tensors(batch, seq,
                    hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype=dtype, group_num=8)
            torch_out = torch_ffn(input, quant_w1, scale1, bias1, quant_w2, scale2, bias2, is_gated)
            tmo_out = tmo_weight_only_group_quant_ffn(input, quant_w1, scale1.to(dtype=input.dtype), bias1,
                    quant_w2, scale2.to(dtype=input.dtype), bias2, quant_bit)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                    0.05, use_MSE=True, use_RAE=True)

    def test_pertoken_smooth_quant_ffn(self):
        for dtype in dtype_list:
            print("test_pertoken_smooth_quant_ffn...")
            batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype  = 3, 5, 512, 768, 8, True, True, dtype
            input, quant_w1, scale1, quant_w2, scale2, bias1, bias2, input_smooth, act_smooth = \
                init_tensors(batch, seq, hidden_size, inner_size, quant_bit, is_gated, with_smooth, dtype)
            torch_out = torch_smooth_quant_ffn(input, quant_w1, scale1, bias1, quant_w2, scale2, bias2,
                                            input_smooth, act_smooth, is_gated)
            tmo_out = tmo_pertoken_smooth_quant_gated_ffn(input, quant_w1, scale1, bias1, quant_w2,
                scale2, bias2, input_smooth, act_smooth, dtype)
            self.assertTensorsEqual(torch_out.cpu().float(), tmo_out.cpu().float(),
                                    0.05, use_MSE=True, use_RAE=True)

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

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