# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/enc_dec/helper.py

import typing
from typing import Union

import numpy as np
import torch  # pytype: disable=import-error

from tensorrt_llm._utils import str_dtype_to_torch


def split(v: Union[np.ndarray, torch.Tensor],
          tp_size: int,
          tp_rank: int,
          dim=0):
    if tp_size == 1:
        if isinstance(v, np.ndarray):
            return np.ascontiguousarray(v.copy())
        else:
            return v.clone().detach()
    assert len(v.shape) > 1 or dim == 0
    if isinstance(v, np.ndarray):
        return np.ascontiguousarray(
            np.split(v, tp_size, axis=dim)[tp_rank].copy())
    else:
        assert v.shape[dim] % tp_size == 0, \
            'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.'
        split_size = v.shape[dim] // tp_size
        return v.split(split_size, dim=dim)[tp_rank].clone().detach()


def reshape(v: torch.Tensor, shape=None):
    if shape is None:
        return v.contiguous()
    else:
        return v.reshape(shape).contiguous()


def fuse_qkv_one_layer(params, attn_module_name, trtllm_layer_name, tp_size,
                       tp_rank, model_type, weight_shape, bias_shape):

    qkv_module_names = get_qkv_module_name(model_type)

    weight = {}

    # fuse weights of q, k, v
    q_w = params[f'{attn_module_name}.{qkv_module_names["q"]}.weight']
    k_w = params[f'{attn_module_name}.{qkv_module_names["k"]}.weight']
    v_w = params[f'{attn_module_name}.{qkv_module_names["v"]}.weight']

    # fuse qkv weight
    shape = q_w.shape  # (do, din)
    qkv_w = torch.cat([q_w, k_w, v_w],
                      dim=0).reshape([3, shape[0], shape[1]])  # (3, do, din)
    qkv_w = split(qkv_w, tp_size, tp_rank, dim=1)
    weight[f'{trtllm_layer_name}.qkv.weight'] = reshape(qkv_w,
                                                        shape=weight_shape)

    # fuse qkv biases if present
    if f'{attn_module_name}.{qkv_module_names["q"]}.bias' in params.keys(
    ) and params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] is not None:
        q_b = params[f'{attn_module_name}.{qkv_module_names["q"]}.bias']
        k_b = params[f'{attn_module_name}.{qkv_module_names["k"]}.bias']
        v_b = params[f'{attn_module_name}.{qkv_module_names["v"]}.bias']
        shape = q_b.shape[0]  # (do,)
        qkv_b = torch.cat([q_b, k_b, v_b], dim=0).reshape([3, shape])  # (3, do)
        qkv_b = split(qkv_b, tp_size, tp_rank, dim=1)
        weight[f'{trtllm_layer_name}.qkv.bias'] = reshape(qkv_b,
                                                          shape=bias_shape)
    return weight


def get_qkv_module_name(model_type):
    if model_type in ["t5", "blip2"]:
        q = "q"
        k = "k"
        v = "v"
    elif model_type == "bart" or model_type == "nmt":
        q = "q_proj"
        k = "k_proj"
        v = "v_proj"
    elif model_type == "pix2struct":
        q = "query"
        k = "key"
        v = "value"
    return {"q": q, "k": k, "v": v}


def convert_weight_to_dtype(params: typing.Dict[str, torch.Tensor],
                            dtype: typing.Optional[np.dtype] = None):
    if dtype is not None:
        assert isinstance(dtype,
                          str), f"dtype must be str, but get type {type(dtype)}"
        for name in params.keys():
            params[name] = params[name].to(str_dtype_to_torch(dtype))