from collections.abc import Iterable
import io
from typing import Any, Dict, List, Tuple, Union, Optional

import pytest
import torch
import transformer_engine as te

import transformer_engine.common.recipe
import transformer_engine.pytorch as te
from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
from transformer_engine.pytorch.tensor.float8_tensor import (
    Float8Quantizer,
    Float8Tensor,
    Float8CurrentScalingQuantizer,
)
from transformer_engine.pytorch.constants import TE_DType, TE_DType_To_Torch
from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported
import transformer_engine_torch as tex

from references.ref_per_tensor_cs import ref_per_tensor_cs_cast

from transformer_engine.pytorch.triton.per_token_group_quant import per_token_quant_int8, per_token_quant_int8_v2, per_token_quant_fp8_to_int8, per_token_quant_fp8_to_int8_v2, channelwise_dequantize, channelwise_dequantize_transA, channelwise_dequantize_transB, per_token_quant_fp8_to_int8_opt
import time



# TN
m = 4096
k = 4096
n = 4096
seed = 0
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
device = "cuda"
out_dtype = torch.int32

# Allocate cuBLAS workspace
workspace_size = 128
workspace = torch.empty(128, dtype=torch.uint8, device=device)
out_quantizer = None
accumulate = False
use_gelu = False
use_bias = False
bias = None
use_grad = False
assert not (use_gelu and use_bias), "Bias and GELU not supported by GEMM"
aux_tensor = torch.empty((m, n), dtype=out_dtype, device=device) if use_gelu else None
out = torch.empty((m, n), dtype=out_dtype, device=device) if accumulate else None
bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype]
use_split_accumulator = False

# bf16 to int8

# transa = True
# transb = False

# x_bf16 = (torch.randn((m, k), device=device)).to(dtype=torch.bfloat16)
# w_bf16 = (torch.randn((n, k), device=device)).to(dtype=torch.bfloat16)

# bf16_out = torch.matmul(x_bf16, w_bf16.t())

# x_int8, x_scales = per_token_quant_int8(x_bf16)
# w_int8, w_scales = per_token_quant_int8(w_bf16)
# # print("x_int8: ", x_int8)
# # print("w_int8: ", w_int8)

# # cuBLAS GEMM
# # return type is out, bias_grad, gelu_input, extra_output
# # We are just capturing out.
# y_int32 = tex.generic_gemm(
#     w_int8,
#     transa,
#     x_int8,
#     transb,
#     out,
#     out_quantizer,
#     TE_DType[out_dtype],
#     bias,
#     bias_dtype,
#     use_gelu,
#     aux_tensor,
#     use_grad,
#     workspace,
#     workspace.shape[0],
#     accumulate,
#     use_split_accumulator,
# )[0]

# # y_int32 = torch._int_mm(x_int8, w_int8.t())
# # print("y_int32: ", y_int32)

# output = channelwise_dequantize_transB(x_scales, w_scales, y_int32)
# # print("out_scales.shape: ", out_scales.shape)
# # print("out_scales: ", out_scales)
# # print("bf16_out: ", bf16_out)
# # print("output: ", output)

# # NN

# transa = False
# transb = False

# dy_bf16 = (torch.randn((m, n), device=device)).to(dtype=torch.bfloat16)
# w_bf16 = (torch.randn((n, k), device=device)).to(dtype=torch.bfloat16)

# bf16_dx = torch.matmul(dy_bf16, w_bf16)

# dy_int8, dy_scales = per_token_quant_int8(dy_bf16)
# w_int8, w_scales = per_token_quant_int8_v2(w_bf16)
# # print("dy_scales.shape: ", dy_scales.shape)
# # print("w_scales.shape: ", w_scales.shape)
# # print("dy_int8: ", dy_int8)
# # print("w_int8: ", w_int8)

# # cuBLAS GEMM
# # return type is out, bias_grad, gelu_input, extra_output
# # We are just capturing out.
# dx_int32 = tex.generic_gemm(
#     w_int8,
#     transa,
#     dy_int8,
#     transb,
#     out,
#     out_quantizer,
#     TE_DType[out_dtype],
#     bias,
#     bias_dtype,
#     use_gelu,
#     aux_tensor,
#     use_grad,
#     workspace,
#     workspace.shape[0],
#     accumulate,
#     use_split_accumulator,
# )[0]

# # dx_int32 = torch._int_mm(dy_int8, w_int8)
# # print("dx_int32: ", dx_int32)

# dx = channelwise_dequantize(dy_scales, w_scales, dx_int32)
# # print("dx_scales.shape: ", dx_scales.shape)
# # print("dx_scales: ", dx_scales)
# # print("bf16_dx: ", bf16_dx)
# # print("dx: ", dx)



# # NT

# transa = False
# transb = True

# dy_bf16 = (torch.randn((m, n), device=device)).to(dtype=torch.bfloat16)
# x_bf16 = (torch.randn((m, k), device=device)).to(dtype=torch.bfloat16)

# bf16_dw = torch.matmul(dy_bf16.t(), x_bf16)

# dy_int8, dy_scales = per_token_quant_int8_v2(dy_bf16)
# x_int8, x_scales = per_token_quant_int8_v2(x_bf16)

# # cuBLAS GEMM
# # return type is out, bias_grad, gelu_input, extra_output
# # We are just capturing out.
# dw_int32 = tex.generic_gemm(
#     x_int8,
#     transa,
#     dy_int8,
#     transb,
#     out,
#     out_quantizer,
#     TE_DType[out_dtype],
#     bias,
#     bias_dtype,
#     use_gelu,
#     aux_tensor,
#     use_grad,
#     workspace,
#     workspace.shape[0],
#     accumulate,
#     use_split_accumulator,
# )[0]

# dw = channelwise_dequantize_transA(dy_scales, x_scales, dw_int32)
# # print("bf16_dw: ", bf16_dw)
# # print("dw: ", dw)



# fp8 to int8


quantizer_e5m2 = Float8CurrentScalingQuantizer(
    fp8_dtype=tex.DType.kFloat8E5M2,
    device="cuda",
    force_pow_2_scales=False,
    amax_epsilon=0.0,
)

quantizer_e4m3 = Float8CurrentScalingQuantizer(
    fp8_dtype=tex.DType.kFloat8E4M3,
    device="cuda",
    force_pow_2_scales=False,
    amax_epsilon=0.0,
)

# current scaling
def to_float8_CS(
    tensor: torch.Tensor,
    fp8_dtype: tex.DType = tex.DType.kFloat8E5M2,
    return_transpose: bool = False,
    force_pow_2_scales: bool = False,
    amax_epsilon: float = 0.0,
) -> Float8Tensor:
    """Cast tensor to FP8"""
    quantizer = quantizer_e5m2 if fp8_dtype == tex.DType.kFloat8E5M2 else quantizer_e4m3
    if return_transpose:
        quantizer.set_usage(rowwise=True, columnwise=True)
    else:
        quantizer.set_usage(rowwise=True, columnwise=False)
    return quantizer(tensor)

# TN

transa = True
transb = False

x_bf16 = (torch.randn((m, k), device=device)).to(dtype=torch.bfloat16)
w_bf16 = (torch.randn((n, k), device=device)).to(dtype=torch.bfloat16)

bf16_out = torch.matmul(x_bf16, w_bf16.t())

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    bf16_out = torch.matmul(x_bf16, w_bf16.t())
torch.cuda.synchronize()
end = time.time()

# x_int8, x_scales = per_token_quant_int8(x_bf16)
# w_int8, w_scales = per_token_quant_int8(w_bf16)
# print("x_int8: ", x_int8)
# print("w_int8: ", w_int8)

# Cast to FP8 and back
x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E4M3)
w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E4M3)
# print("x_fp8: ", x_fp8._data.view(dtype=torch.float8_e4m3fn))
# print("w_fp8: ", w_fp8._data.view(dtype=torch.float8_e4m3fn))

x_int8, x_scales = per_token_quant_fp8_to_int8(x_fp8._data.view(dtype=torch.float8_e4m3fn), x_fp8._scale_inv, False)
w_int8, w_scales = per_token_quant_fp8_to_int8(w_fp8._data.view(dtype=torch.float8_e4m3fn), w_fp8._scale_inv, False)
# print("x_int8: ", x_int8)
# print("w_int8: ", w_int8)

# cuBLAS GEMM
# return type is out, bias_grad, gelu_input, extra_output
# We are just capturing out.
y_int32 = tex.generic_gemm(
    w_int8,
    transa,
    x_int8,
    transb,
    out,
    out_quantizer,
    TE_DType[out_dtype],
    bias,
    bias_dtype,
    use_gelu,
    aux_tensor,
    use_grad,
    workspace,
    workspace.shape[0],
    accumulate,
    use_split_accumulator,
)[0]

# y_int32 = torch._int_mm(x_int8, w_int8.t())
# print("y_int32: ", y_int32)

output = channelwise_dequantize_transB(x_scales, w_scales, y_int32)
# print("out_scales.shape: ", out_scales.shape)
# print("out_scales: ", out_scales)
print("bf16_out: ", bf16_out)
print("output: ", output)

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E4M3)
    # w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E4M3)
    x_int8, x_scales = per_token_quant_fp8_to_int8(x_fp8._data.view(dtype=torch.float8_e4m3fn), x_fp8._scale_inv, False)
    w_int8, w_scales = per_token_quant_fp8_to_int8(w_fp8._data.view(dtype=torch.float8_e4m3fn), w_fp8._scale_inv, False)
    y_int32 = tex.generic_gemm(
        w_int8,
        transa,
        x_int8,
        transb,
        out,
        out_quantizer,
        TE_DType[out_dtype],
        bias,
        bias_dtype,
        use_gelu,
        aux_tensor,
        use_grad,
        workspace,
        workspace.shape[0],
        accumulate,
        use_split_accumulator,
    )[0]
    output = channelwise_dequantize_transB(x_scales, w_scales, y_int32)
torch.cuda.synchronize()
end = time.time()

# NN

# transa = True
transa = False
transb = False

dy_bf16 = (torch.randn((m, n), device=device)).to(dtype=torch.bfloat16)
w_bf16 = (torch.randn((n, k), device=device)).to(dtype=torch.bfloat16)

bf16_dx = torch.matmul(dy_bf16, w_bf16)

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    bf16_dx = torch.matmul(dy_bf16, w_bf16)
torch.cuda.synchronize()
end = time.time()

# Cast to FP8 and back
dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
# w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)

dy_int8, dy_scales = per_token_quant_fp8_to_int8(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
w_int8, w_scales = per_token_quant_fp8_to_int8_opt(w_fp8._data.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)
# w_int8, w_scales = per_token_quant_fp8_to_int8_v2(w_fp8._data.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)
# w_int8, w_scales = per_token_quant_fp8_to_int8(w_fp8._transpose.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)

# print("dy_scales.shape: ", dy_scales.shape)
# print("w_scales.shape: ", w_scales.shape)
# print("dy_int8: ", dy_int8)
# print("w_int8: ", w_int8)
# print("w_scales: ", w_scales)

# cuBLAS GEMM
# return type is out, bias_grad, gelu_input, extra_output
# We are just capturing out.
dx_int32 = tex.generic_gemm(
    w_int8,
    transa,
    dy_int8,
    transb,
    out,
    out_quantizer,
    TE_DType[out_dtype],
    bias,
    bias_dtype,
    use_gelu,
    aux_tensor,
    use_grad,
    workspace,
    workspace.shape[0],
    accumulate,
    use_split_accumulator,
)[0]

# dx_int32 = torch._int_mm(dy_int8, w_int8)
# print("dx_int32: ", dx_int32)

dx = channelwise_dequantize(dy_scales, w_scales, dx_int32)
# dx = channelwise_dequantize_transB(dy_scales, w_scales, dx_int32)

# print("dx_scales.shape: ", dx_scales.shape)
# print("dx_scales: ", dx_scales)
print("bf16_dx: ", bf16_dx)
print("dx: ", dx)

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
    # w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
    # w_fp8 = to_float8_CS(w_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)
    dy_int8, dy_scales = per_token_quant_fp8_to_int8(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
    w_int8, w_scales = per_token_quant_fp8_to_int8_opt(w_fp8._data.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)
    # w_int8, w_scales = per_token_quant_fp8_to_int8_v2(w_fp8._data.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)
    # w_int8, w_scales = per_token_quant_fp8_to_int8(w_fp8._transpose.view(dtype=torch.float8_e5m2), w_fp8._scale_inv, False)
    dx_int32 = tex.generic_gemm(
        w_int8,
        transa,
        dy_int8,
        transb,
        out,
        out_quantizer,
        TE_DType[out_dtype],
        bias,
        bias_dtype,
        use_gelu,
        aux_tensor,
        use_grad,
        workspace,
        workspace.shape[0],
        accumulate,
        use_split_accumulator,
    )[0]
    dx = channelwise_dequantize(dy_scales, w_scales, dx_int32)
    # dx = channelwise_dequantize_transB(dy_scales, w_scales, dx_int32)
torch.cuda.synchronize()
end = time.time()


# NT

# transa = True
# transb = False
transa = False
transb = True

dy_bf16 = (torch.randn((m, n), device=device)).to(dtype=torch.bfloat16)
x_bf16 = (torch.randn((m, k), device=device)).to(dtype=torch.bfloat16)
seed = 0
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

bf16_dw = torch.matmul(dy_bf16.t(), x_bf16)

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    bf16_dw = torch.matmul(dy_bf16.t(), x_bf16)
torch.cuda.synchronize()
end = time.time()

# Cast to FP8 and back
# dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)
# x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)
dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E5M2)

# dy_int8, dy_scales = per_token_quant_fp8_to_int8_v2(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
# x_int8, x_scales = per_token_quant_fp8_to_int8_v2(x_fp8._data.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)
dy_int8, dy_scales = per_token_quant_fp8_to_int8_opt(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
x_int8, x_scales = per_token_quant_fp8_to_int8_opt(x_fp8._data.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)
# dy_int8, dy_scales = per_token_quant_fp8_to_int8(dy_fp8._transpose.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
# x_int8, x_scales = per_token_quant_fp8_to_int8(x_fp8._transpose.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)

# cuBLAS GEMM
# return type is out, bias_grad, gelu_input, extra_output
# We are just capturing out.
dw_int32 = tex.generic_gemm(
    x_int8,
    transa,
    dy_int8,
    transb,
    out,
    out_quantizer,
    TE_DType[out_dtype],
    bias,
    bias_dtype,
    use_gelu,
    aux_tensor,
    use_grad,
    workspace,
    workspace.shape[0],
    accumulate,
    use_split_accumulator,
)[0]

dw = channelwise_dequantize_transA(dy_scales, x_scales, dw_int32)
# dw = channelwise_dequantize_transB(dy_scales, x_scales, dw_int32)
print("bf16_dw: ", bf16_dw)
print("dw: ", dw)

torch.cuda.synchronize()
start = time.time()
for i in range(20):
    # dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)
    # x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E5M2, return_transpose=True)
    # dy_fp8 = to_float8_CS(dy_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
    # x_fp8 = to_float8_CS(x_bf16, fp8_dtype=tex.DType.kFloat8E5M2)
    
    # dy_int8, dy_scales = per_token_quant_fp8_to_int8_v2(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
    # x_int8, x_scales = per_token_quant_fp8_to_int8_v2(x_fp8._data.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)
    dy_int8, dy_scales = per_token_quant_fp8_to_int8_opt(dy_fp8._data.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
    x_int8, x_scales = per_token_quant_fp8_to_int8_opt(x_fp8._data.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)
    # dy_int8, dy_scales = per_token_quant_fp8_to_int8(dy_fp8._transpose.view(dtype=torch.float8_e5m2), dy_fp8._scale_inv, False)
    # x_int8, x_scales = per_token_quant_fp8_to_int8(x_fp8._transpose.view(dtype=torch.float8_e5m2), x_fp8._scale_inv, False)
    dw_int32 = tex.generic_gemm(
        x_int8,
        transa,
        dy_int8,
        transb,
        out,
        out_quantizer,
        TE_DType[out_dtype],
        bias,
        bias_dtype,
        use_gelu,
        aux_tensor,
        use_grad,
        workspace,
        workspace.shape[0],
        accumulate,
        use_split_accumulator,
    )[0]
    dw = channelwise_dequantize_transA(dy_scales, x_scales, dw_int32)
    # dw = channelwise_dequantize_transB(dy_scales, x_scales, dw_int32)
torch.cuda.synchronize()
end = time.time()



# bacth gemm wgrad

m = 32
k = 32
n = 32
b = 4

transa = False
transb = True

dy_int8 = (torch.randn((b, m, n), device=device)).to(dtype=torch.int8)
x_int8 = (torch.randn((b, m, k), device=device)).to(dtype=torch.int8)
seed = 0
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

int32_dw_list = []
for i in range(b):
    int32_dw = torch._int_mm(dy_int8[i].t(), x_int8[i])
    # bf16_dw = torch.matmul(dy_int8[i].t(), x_int8[i])
    int32_dw_list.append(int32_dw)

batched_int32_dw = torch.stack(int32_dw_list)
# print("batched_int32_dw.shape: ", batched_int32_dw.shape)
# print("batched_int32_dw: ", batched_int32_dw)

out_dtype = torch.int32
out = torch.empty((b, n, k), dtype=out_dtype, device=device)

te_dw = tex.generic_batchgemm(
    x_int8.view(-1, x_int8.size(-1)),
    transa,
    dy_int8.view(-1, dy_int8.size(-1)),
    transb,
    out.view(-1, out.size(-1)),
    b,
    out_quantizer,
    TE_DType[out_dtype],
    bias,
    bias_dtype,
    use_gelu,
    aux_tensor,
    use_grad,
    workspace,
    workspace.shape[0],
    accumulate,
    use_split_accumulator,
)[0]
# print("te_dw.shape: ", te_dw.view(b, -1, te_dw.size(-1)).shape)
# print("te_dw: ", te_dw.view(b, -1, te_dw.size(-1)))
torch.testing.assert_close(te_dw.view(b, -1, te_dw.size(-1)), batched_int32_dw, atol=0, rtol=0)
