from typing import Literal, Dict

# Implementation asm for fp4 to bf16, using twiddling
# Reference: https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/tensor_details/layout_details/hopper_value.py#L11-L18
decode_f4_to_bf16_twiddling = """
// N should be the number of elements processed by one thread
template<typename T1, typename T2>
__device__ void decode_fp4_to_bf16_twiddling(T1 *B_local, T2 *B_local_decode, const int N = 8) {
  #pragma unroll
  for (int i = 0; i < N; ++i) {
    uint B_dequantize_local_vec[4];
    uint tmp, bias, d0, d1, d2, d3, d4, d5, d6;
    asm volatile(
      // To handle the endianness issue
      "prmt.b32 %13, %4, 0, 0x0123;"
      "mov.b32 %12, 0x7e807e80;"
      "and.b32 %0, %13, 0b10000001110000001000000111000000;"
      "mul.bf16x2 %0, %0, %12;"
      "shl.b32 %1, %13, 3;"
      "and.b32 %1, %1, 0b10000001110000001000000111000000;"          
      "mul.bf16x2 %1, %1, %12;"
      "shl.b32 %2, %13, 6;"
      "and.b32 %2, %2, 0b10000001110000001000000111000000;"
      "mul.bf16x2 %2, %2, %12;"
      "shl.b32 %5, %13, 1;"
      "and.b32 %6, %5, 0b10000000000000001000000000000000;"
      "shr.b32 %7, %13, 3;"
      "and.b32 %8, %7, 0b00000001100000000000000110000000;"
      "or.b32 %9, %6, %8;"
      "shr.b32 %10, %13, 7;"
      "and.b32 %11, %10, 0b00000000010000000000000001000000;"
      "or.b32 %3, %9, %11;"
      "mul.bf16x2 %3, %3, %12;"
      :"=r"(B_dequantize_local_vec[0])
      ,"=r"(B_dequantize_local_vec[1])
      ,"=r"(B_dequantize_local_vec[2])
      ,"=r"(B_dequantize_local_vec[3])
      :"r"(*(uint*)&B_local[i << 2]), "r"(d0), "r"(d1), "r"(d2), "r"(d3), "r"(d4), "r"(d5), "r"(d6), "r"(bias), "r"(tmp)
    );
    for (int j = 0; j < 4; ++j) {
      // Pay attention to the big-endianness issue
      B_local_decode[(i << 3) + j] = reinterpret_cast<T2*>(&B_dequantize_local_vec[j])[1];
      B_local_decode[(i << 3) + j + 4] = reinterpret_cast<T2*>(&B_dequantize_local_vec[j])[0];
    }    
  }
  // Check if the synchronization is needed
}
"""


def get_mxfp_intrin_group(
    out_dtype: Literal["float16", "bfloat16"] = "bfloat16",
    source_format: Literal["int", "uint"] = "uint",
    source_bit: int = 4,
    storage_dtype: Literal["int32", "int8", "uint8"] = "uint8",
    use_twiddling: bool = False,
) -> Dict[str, str]:
    """
    This function is used to get the intrinsic group of the MXFP operation to avoid the overhead of fast decoding.
    MXFP is a type of logic operation that takes three inputs. The intrinsic group refers to the set of
    intrinsic operations that can be performed on these inputs. This function retrieves and returns this group.
    """
    assert out_dtype in ["float16", "bfloat16"
                        ], f"Invalid out_dtype: {out_dtype}. Expected 'float16' or 'bfloat16'."
    assert source_format in ["int", "uint"
                            ], f"Invalid source_format: {source_format}. Expected 'int' or 'uint'."
    assert storage_dtype in [
        "int32", "int8", "uint8"
    ], f"Invalid storage_dtype: {storage_dtype}. Expected 'int32' or 'int8' or 'uint8'."

    dtype_map = {"float16": "f16", "bfloat16": "bf16"}
    key = f"fp{source_bit}_to_{dtype_map[out_dtype]}"
    if use_twiddling:
        key += "_twiddling"

    import_c_map = {
        "fp4_to_bf16_twiddling": decode_f4_to_bf16_twiddling,
    }

    func_name = f"decode_fp{source_bit}_to_{dtype_map[out_dtype]}"
    if use_twiddling:
        func_name += "_twiddling"

    return {
        "func_name": func_name,
        "c_source": import_c_map[key],
    }
