import sys
import os
import importlib.util
import argparse

patch_select = os.environ.get('MODEL_PATCH', '').lower()
if 'deepseek' in patch_select:
    custom_modeling_path = os.path.abspath("./model/lade/models/modeling_llama_kv_ds.py")
    module_name = "transformers.models.llama.modeling_llama"
    spec = importlib.util.spec_from_file_location(module_name, custom_modeling_path)
    custom_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(custom_module)
    sys.modules[module_name] = custom_module

elif 'qwen3' in patch_select:
    custom_modeling_path = os.path.abspath("./model/lade/models/modeling_qwen3.py")
    module_name = "transformers.models.qwen3.modeling_qwen3"
    spec = importlib.util.spec_from_file_location(module_name, custom_modeling_path)
    custom_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(custom_module)
    sys.modules[module_name] = custom_module

import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastchat.utils import str_to_torch_dtype
from evaluation.eval import run_eval, reorg_answer_file
from model.lade.utils import augment_all, config_lade
from model.lade.decoding import CONFIG_MAP

def lookahead_forward(inputs, model, tokenizer, max_new_tokens):
    input_ids = inputs.input_ids
    output_ids, idx, accept_length_list = model.generate(
        torch.as_tensor(input_ids).cuda(),
        do_sample=False,
        temperature=0.0,
        max_new_tokens=max_new_tokens,
    )
    new_token = len(output_ids[0][len(input_ids[0]):])
    return output_ids, new_token, idx+1, accept_length_list


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, required=True, help="The path to the weights. This can be a local folder or a Hugging Face repo ID.")
    parser.add_argument("--model-id", type=str, required=True, help="A custom name for the model.")
    parser.add_argument("--bench-name", type=str, default="mt_bench", help="The name of the benchmark question set.")
    parser.add_argument("--question-begin", type=int, help="A debug option. The begin index of questions.")
    parser.add_argument("--question-end", type=int, help="A debug option. The end index of questions.")
    parser.add_argument("--answer-file", type=str, help="The output answer file.")
    parser.add_argument("--max-new-tokens", type=int, default=1024, help="The maximum number of new generated tokens.")
    parser.add_argument("--num-choices", type=int, default=1, help="How many completion choices to generate.")
    parser.add_argument("--num-gpus-per-model", type=int, default=1, help="The number of GPUs per model.")
    parser.add_argument("--num-gpus-total", type=int, default=1, help="The total number of GPUs.")
    parser.add_argument("--level", type=int, default=3)
    parser.add_argument("--window", type=int, default=10)
    parser.add_argument("--guess", type=int, default=10)
    parser.add_argument("--dtype", type=str, default="float16", choices=["float32", "float64", "float16", "bfloat16"], help="Override the default dtype. If not set, it will use float16 on GPU.")
    parser.add_argument("--use-cot-data", action="store_true", help="Use cot data. If not set, will use False by default.")
    parser.add_argument("--think-twice", action="store_true", help="Use original deepseek forward(). If not set, will use False by default.")
    parser.add_argument("--BON", action="store_true", help="Use original deepseek forward(). If not set, will use False by default.")
    args = parser.parse_args()
    
    if int(os.environ.get("USE_LADE", 0)):
        augment_all(args.model_path)
        config_lade(LEVEL=args.level, WINDOW_SIZE=args.window, GUESS_SET_SIZE=args.guess, DEBUG=0,
                         USE_FLASH=0, DIST_WORKERS=len(os.environ.get("CUDA_VISIBLE_DEVICES").split(",")))
        print("lade activated config: ", CONFIG_MAP)

    if args.use_cot_data:
        question_file = f"data/{args.bench_name}/question_cot.jsonl"
    else:
        question_file = f"data/{args.bench_name}/question.jsonl"

    if args.answer_file:
        answer_file = args.answer_file
    else:
        answer_file = f"data/{args.bench_name}/model_answer/{args.model_id}.jsonl"

    print(f"Output to {answer_file}")
    
    _support_cot_models = ["DeepSeek", "Qwen3"]
    cot_model_flag = any(model in args.model_path for model in _support_cot_models)
    
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=str_to_torch_dtype(args.dtype),
        low_cpu_mem_usage=True,
        device_map="auto"
    )

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)

    run_eval(
        model=model,
        tokenizer=tokenizer,
        forward_func=lookahead_forward,
        model_id=args.model_id,
        question_file=question_file,
        question_begin=args.question_begin,
        question_end=args.question_end,
        answer_file=answer_file,
        max_new_tokens=args.max_new_tokens,
        cot_model_flag=cot_model_flag,
        think_twice=args.think_twice,
        BON=args.BON,
        num_choices=args.num_choices,
        num_gpus_per_model=args.num_gpus_per_model,
        num_gpus_total=args.num_gpus_total,
    )

    reorg_answer_file(answer_file)
