import argparse
import json
import os
from typing import Dict, List

import torch

from rtp_llm.model_factory import ModelFactory
from rtp_llm.tools.fake_util import copy_from_model, generate_fake_model
from rtp_llm.utils.model_weight import W
from rtp_llm.utils.util import load_ckpt


def default_save_config_func(
    model_type,
    dest_path: str,
    layer: int,
    head: int,
    head_kv: int,
    head_size: int,
    ffn_hidden_size: int,
    ffn_inter_padding_size: int,
    vocab_size: int,
):
    config = {
        "model_type": model_type,
        "add_bias_linear": False,
        "add_qkv_bias": True,
        "apply_query_key_layer_scaling": True,
        "apply_residual_connection_post_layernorm": False,
        "attention_dropout": 0.0,
        "attention_softmax_in_fp32": True,
        "bias_dropout_fusion": True,
        "ffn_hidden_size": ffn_hidden_size,
        "ffn_inter_padding_size": ffn_inter_padding_size,
        "fp32_residual_connection": False,
        "hidden_dropout": 0.0,
        "hidden_size": head * head_size,
        "kv_channels": head_size,
        "layernorm_epsilon": 1e-05,
        "multi_query_attention": True,
        "multi_query_group_num": head_kv,
        "num_attention_heads": head,
        "num_layers": layer,
        "original_rope": True,
        "padded_vocab_size": vocab_size,
        "post_layer_norm": True,
        "rmsnorm": True,
        "seq_length": 32768,
        "use_cache": True,
        "torch_dtype": "float16",
        "tie_word_embeddings": False,
        "eos_token_id": 2,
        "pad_token_id": 0,
    }
    # save to config.json
    json.dump(config, open(os.path.join(dest_path, "config.json"), "w"), indent=2)


def fake_model_impl(
    model_type: str,
    save_config_func,
    post_rewrite_func,
    dest_path: str,
    layer_num: int,
    head_num: int,
    head_kv_num: int,
    head_size: int,
    ffn_hidden_size: int,
    ffn_inter_padding_size: int,
    ffn_gate_active: bool,
    ffn_w1_w3_independ: bool,
    vocab_size: int,
    input_model: str | None,
):
    hidden_size = head_num * head_size
    qkv_hidden_size = hidden_size + head_kv_num * head_size * 2
    new_params: Dict[str, torch.Tensor] = {}

    model_weight_info = ModelFactory.get_weight_cls(model_type)(
        hidden_size=hidden_size,
        inter_size=ffn_hidden_size * 2,
        num_heads=head_num,
        num_heads_kv=head_kv_num,
        tp_size=1,
        int8_mode=0,
        num_layers=layer_num,
    )

    model_weight_info._lm_head = False
    model_weight_info._transformer_prefix = False
    weight_info = model_weight_info.get_weight_info()
    print(weight_info)

    weight_dict = {
        W.embedding: "",
        W.lm_head: "",
        W.pre_decoder_ln_gamma: "",
        W.pre_decoder_ln_beta: "",
        W.final_ln_gamma: "",
        W.final_ln_beta: "",
    }

    for weight_item in weight_info.weights:
        if weight_item.name in weight_dict and len(weight_item.weights) > 0:
            weight_dict[weight_item.name] = weight_item.weights[0].name

    shape_map: Dict[str, List[int]] = {
        weight_dict[W.embedding]: [vocab_size, hidden_size],
        weight_dict[W.lm_head]: [vocab_size, hidden_size],
        weight_dict[W.pre_decoder_ln_gamma]: [hidden_size],
        weight_dict[W.pre_decoder_ln_beta]: [hidden_size],
        weight_dict[W.final_ln_gamma]: [hidden_size],
        weight_dict[W.final_ln_beta]: [hidden_size],
    }

    layer_weight_dict = {
        W.pre_ln_gamma: "",
        W.pre_ln_beta: "",
        W.attn_qkv_w: "",
        W.attn_qkv_b: "",
        W.attn_ln_gamma: "",
        W.attn_ln_beta: "",
        W.attn_o_w: "",
        W.attn_o_b: "",
        W.ffn_w1: "",
        W.ffn_b1: "",
        W.ffn_ln_gamma: "",
        W.ffn_ln_beta: "",
        W.ffn_w2: "",
        W.ffn_b2: "",
        W.ffn_w3: "",
        W.ffn_b3: "",
        W.post_ln_gamma: "",
        W.post_ln_beta: "",
    }

    for layer_weight_item in weight_info.layer_weights:
        if (
            layer_weight_item.name in layer_weight_dict
            and len(layer_weight_item.weights) > 0
        ):
            layer_weight_dict[layer_weight_item.name] = layer_weight_item.weights[
                0
            ].name

    for i in range(layer_num):
        shape_map[layer_weight_dict[W.pre_ln_gamma].format(i=str(i))] = [hidden_size]
        shape_map[layer_weight_dict[W.pre_ln_beta].format(i=str(i))] = [hidden_size]
        shape_map[layer_weight_dict[W.attn_qkv_w].format(i=str(i))] = [
            qkv_hidden_size,
            hidden_size,
        ]
        shape_map[layer_weight_dict[W.attn_qkv_b].format(i=str(i))] = [qkv_hidden_size]
        shape_map[layer_weight_dict[W.attn_ln_gamma].format(i=str(i))] = [hidden_size]
        shape_map[layer_weight_dict[W.attn_ln_beta].format(i=str(i))] = [hidden_size]
        shape_map[layer_weight_dict[W.attn_o_w].format(i=str(i))] = [
            hidden_size,
            hidden_size,
        ]
        shape_map[layer_weight_dict[W.attn_o_b].format(i=str(i))] = [hidden_size]

        if ffn_gate_active:
            if ffn_w1_w3_independ:
                shape_map[layer_weight_dict[W.ffn_w1].format(i=str(i))] = [
                    ffn_hidden_size,
                    hidden_size,
                ]
                shape_map[layer_weight_dict[W.ffn_b1].format(i=str(i))] = [
                    ffn_hidden_size
                ]
                shape_map[layer_weight_dict[W.ffn_w3].format(i=str(i))] = [
                    ffn_hidden_size,
                    hidden_size,
                ]
                shape_map[layer_weight_dict[W.ffn_b3].format(i=str(i))] = [
                    ffn_hidden_size
                ]
            else:
                shape_map[layer_weight_dict[W.ffn_w1].format(i=str(i))] = [
                    ffn_hidden_size * 2,
                    hidden_size,
                ]
                shape_map[layer_weight_dict[W.ffn_b1].format(i=str(i))] = [
                    ffn_hidden_size * 2
                ]
        else:
            shape_map[layer_weight_dict[W.ffn_w1].format(i=str(i))] = [
                ffn_hidden_size,
                hidden_size,
            ]
            shape_map[layer_weight_dict[W.ffn_b1].format(i=str(i))] = [ffn_hidden_size]

        shape_map[layer_weight_dict[W.ffn_ln_gamma].format(i=str(i))] = [
            ffn_hidden_size
        ]
        shape_map[layer_weight_dict[W.ffn_ln_beta].format(i=str(i))] = [ffn_hidden_size]

        shape_map[layer_weight_dict[W.ffn_w2].format(i=str(i))] = [
            hidden_size,
            ffn_hidden_size,
        ]
        shape_map[layer_weight_dict[W.ffn_b2].format(i=str(i))] = [hidden_size]

        shape_map[layer_weight_dict[W.post_ln_gamma].format(i=str(i))] = [hidden_size]
        shape_map[layer_weight_dict[W.post_ln_beta].format(i=str(i))] = [hidden_size]

    if "" in shape_map:
        del shape_map[""]
    assert len(shape_map) > 0
    print("shape_map = ", shape_map)

    file_name = f"fake_{model_type}_{layer_num}_{head_num}_{head_kv_num}_{head_size}_{ffn_hidden_size}_{vocab_size}"
    if input_model:
        model = load_ckpt(input_model)
        model_weight_info.process_meta(model)
        new_params = copy_from_model(shape_map, model)
        file_name += "_copy"
    else:
        new_params = generate_fake_model(shape_map)

    file_name += ".pt"

    if not os.path.exists(dest_path):
        print(f"{dest_path} not exist, creating...")
        os.makedirs(dest_path)

    # save config
    print("saving config.json...")
    if save_config_func:
        save_config_func(
            model_type,
            dest_path,
            layer_num,
            head_num,
            head_kv_num,
            head_size,
            ffn_hidden_size,
            ffn_inter_padding_size,
            vocab_size,
        )

    # save model
    print("saving model...")
    file_name = os.path.join(dest_path, file_name)
    if post_rewrite_func:
        new_params = post_rewrite_func(new_params)

    print(new_params)

    torch.save(new_params, file_name)

    print(f"save finished, save path: {dest_path}")


class DefaultModelConfig:
    layer_num: int = 0
    head_num: int = 0
    head_kv_num: int = 0
    head_size: int = 0
    ffn_hidden_size: int = 0
    ffn_inter_padding_size: int = 0
    ffn_gate_active: bool = True
    ffn_w1_w3_independ: bool = False
    vocab_size: int = 0


def fake_model(
    model_type: str,
    default_values: DefaultModelConfig,
    save_config_func=None,
    post_rewrite_func=None,
):
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", "-p", help="saved path", required=True)
    parser.add_argument(
        "--layer", "-l", help="layer number", default=default_values.layer_num, type=int
    )
    parser.add_argument(
        "--head", "-d", help="head number", default=default_values.head_num, type=int
    )
    parser.add_argument(
        "--head_kv",
        "-k",
        help="kv head number",
        default=default_values.head_kv_num,
        type=int,
    )
    parser.add_argument(
        "--head_size",
        "-s",
        help="head size",
        default=default_values.head_size,
        type=int,
    )
    parser.add_argument(
        "--ffn_hidden_size",
        "-f",
        help="ffn hidden size",
        default=default_values.ffn_hidden_size,
        type=int,
    )
    parser.add_argument(
        "--ffn_inter_padding_size",
        "-e",
        help="ffn inter padding size",
        default=default_values.ffn_inter_padding_size,
        type=int,
    )
    parser.add_argument(
        "--ffn_gate_active",
        "-g",
        help="ffn gate active",
        default=default_values.ffn_gate_active,
        type=bool,
    )
    parser.add_argument(
        "--ffn_w1_w3_independ",
        "-r",
        help="ffn w1 w3 weight independ",
        default=default_values.ffn_w1_w3_independ,
        type=bool,
    )
    parser.add_argument(
        "--vocab", "-v", help="vocab size", default=default_values.vocab_size, type=int
    )
    parser.add_argument("--input", "-i", help="input model", default=None, type=str)
    args = parser.parse_args()

    fake_model_impl(
        model_type,
        save_config_func,
        post_rewrite_func,
        args.path,
        args.layer,
        args.head,
        args.head_kv,
        args.head_size,
        args.ffn_hidden_size,
        args.ffn_inter_padding_size,
        args.ffn_gate_active,
        args.ffn_w1_w3_independ,
        args.vocab,
        args.input,
    )
