"""Functions for pre-sharding weights"""

import logging
from collections.abc import Sequence
from typing import Any, Callable, Dict, Tuple  # noqa: UP035

from tvm import IRModule, relax
from tvm.relax.frontend import nn
from tvm.runtime import Device, Tensor
from tvm.s_tir import dlight as dl
from tvm.target import Target

logger = logging.getLogger("preshard")


def _sharded_param_name(param_name, worker_id):
    return f"{param_name}_shard-{worker_id}"


def _create_shard_func(bb: relax.BlockBuilder, param: nn.Parameter, tensor_parallel_shards: int):
    shard_strategy = param.attrs.get("shard_strategy", None)
    # generate tirx shard function
    tir_func = shard_strategy.gen_tir(shards=tensor_parallel_shards, weight=param)
    tir_func = tir_func.with_attr("global_symbol", f"{shard_strategy.name}_tir")
    # add tirx shard function to the IRModule
    tir_gvar = bb.add_func(tir_func, func_name=f"{shard_strategy.name}_tir")
    # create relax function that
    #     1. shard weight with tirx shard function, result: [num_shards, *sharded_weight_shape]
    #     2. split the sharded weight along dim 0, result: num_shards * [1, *sharded_weight_shape]
    #     3. squeeze the 0th-dim of all shards, result: num_shards * [*sharded_weight_shape]
    weight_shape = param.shape
    weight_shape[shard_strategy.dim] = weight_shape[shard_strategy.dim] * tensor_parallel_shards
    sharded_weight_shape = [tensor_parallel_shards, *param.shape]
    weight_var = relax.Var("weight", relax.TensorType(weight_shape, param.dtype))
    with bb.function(name=shard_strategy.name, params=[weight_var]):
        with bb.dataflow():
            lv0 = bb.emit(
                relax.call_tir(
                    tir_gvar,
                    weight_var,
                    out_ty=relax.TensorType(sharded_weight_shape, param.dtype),
                )
            )
            lv1 = bb.emit(relax.op.split(lv0, indices_or_sections=tensor_parallel_shards, axis=0))
            output_vars = []
            for i in range(tensor_parallel_shards):
                lvi = bb.emit(relax.TupleGetItem(lv1, i))
                squeezed_lvi = bb.emit(relax.op.squeeze(lvi, 0))
                output_vars.append(squeezed_lvi)
            gv = bb.emit_output(output_vars)
        bb.emit_func_output(gv)


def _compile_shard_funcs(mod: IRModule, device: Device):
    target = Target.from_device(device)
    with target:
        mod = relax.transform.LegalizeOps()(mod)
        mod = dl.ApplyDefaultSchedule(
            dl.gpu.Matmul(),
            dl.gpu.GEMV(),
            dl.gpu.Reduction(),
            dl.gpu.GeneralReduction(),
            dl.gpu.Fallback(),
        )(mod)
    ex = relax.build(mod, target=target)
    vm = relax.VirtualMachine(ex, device)
    return vm


def apply_preshard(
    named_params: Dict[str, nn.Parameter],  # noqa: UP006
    tensor_parallel_shards: int,
    args: Any,
) -> Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]]:  # noqa: UP006
    """Apply pre-sharding to the named parameters.

    Parameters
    ----------
    named_params : Dict[str, nn.Parameter]
        The named parameters of the model. If the model is quantized, the named parameters should
        the state dictionary of the quantized model.
    tensor_parallel_shards : int
        The number of tensor parallel shards.
    args : Any
        The parsed arguments of weight conversion.

    Returns
    -------
    Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]
        The updated named parameters and the mapping from parameter name to the shard function.
    """
    bb = relax.BlockBuilder()
    param_to_shard_func = {}
    shard_func_names = set()
    new_named_params: Dict[str, nn.Parameter] = {}  # noqa: UP006
    has_shard_strategy = False
    for name, param in named_params.items():
        shard_strategy = param.attrs.get("shard_strategy", None)
        if shard_strategy is not None:
            has_shard_strategy = True
            for i in range(tensor_parallel_shards):
                new_named_params[_sharded_param_name(name, i)] = param
            # create shard functions
            param_to_shard_func[name] = shard_strategy.name
            if shard_strategy.name not in shard_func_names:
                _create_shard_func(bb, param, tensor_parallel_shards)
                shard_func_names.add(shard_strategy.name)
        else:
            new_named_params[name] = param

    if not has_shard_strategy:
        logger.warning(
            "No parameters with 'shard_strategy' found."
            "At least one parameter must have a 'shard_strategy' for presharding. "
            "The model will continue to convert weights in a non-presharded manner."
        )

    mod = bb.finalize()
    vm = _compile_shard_funcs(mod, args.device)

    for name in param_to_shard_func:
        param_to_shard_func[name] = vm[param_to_shard_func[name]]
    return new_named_params, param_to_shard_func