import time
import torch
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

# 固定随机种子，保证采样逻辑尽量可比
set_seed(3407)

# ==================== 对决模型路径配置 ====================
MODEL_A_PATH = "./Qwen3-4B-Thinking-2507"  # 1. 原始 36 层完整模型
MODEL_B_PATH = (
    "./QiMing-Polaris-Qwen3-4B-Final-Accelerated"  # 2. 终极 32 层硬剪枝加速版模型
)

test_prompt = "What is the 'Burden-based Training' method?"

# 统一生成参数
gen_kwargs = {
    "max_new_tokens": 256,
    "temperature": 0.7,
    "top_p": 0.8,
    "do_sample": True,
}


def test_model(model_path, model_name):
    print(f"\n==========================================================")
    print(f"🚀 正在测试模型: [{model_name}]")
    print(f"📁 模型路径    : {model_path}")
    print(f"==========================================================")

    # 1. 加载
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True
    )

    # 2. 格式化 Prompt
    messages = [{"role": "user", "content": test_prompt}]
    formatted_input = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(formatted_input, return_tensors="pt").to("cuda")

    # 3. GPU 预热
    print("🔥 正在预热 GPU...")
    with torch.no_grad():
        _ = model.generate(**inputs, max_new_tokens=10)

    # 4. 正式计时生成
    print("⚡ 开始生成...")
    torch.cuda.synchronize()
    start_time = time.time()

    with torch.no_grad():
        outputs = model.generate(
            **inputs, pad_token_id=tokenizer.eos_token_id, **gen_kwargs
        )

    torch.cuda.synchronize()
    latency = time.time() - start_time

    # 5. 解码与数据统计
    generated_tokens = outputs[0][inputs.input_ids.shape[1] :]
    response_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
    num_tokens = len(generated_tokens)
    tps = num_tokens / latency if latency > 0 else 0

    print(f"\n💬 [{model_name}] 回答:\n{response_text}")
    print("-" * 58)
    print(
        f"⏱️ 耗时: {latency:.4f} 秒 | 生成 Token 数: {num_tokens} | 吞吐率: {tps:.2f} Tokens/s"
    )

    # 6. 清理 GPU 显存
    del model, tokenizer
    gc.collect()
    torch.cuda.empty_cache()

    return {
        "name": model_name,
        "latency": latency,
        "num_tokens": num_tokens,
        "tps": tps,
    }


print("==========================================================")
print("⚔️ 启动终极对决: 原始 36 层模型 VS 终极 32 层加速模型")
print("==========================================================")

# 依次测试两个成品模型
res_A = test_model(MODEL_A_PATH, "原始 36 层模型 (Base)")
res_B = test_model(MODEL_B_PATH, "终极 32 层加速模型 (Final-Accelerated)")

# 终极对决结算表
print("\n" + "=" * 60)
print("🏆 最终成品模型对决结算表 (Final Models Duel)")
print("=" * 60)
print(
    f"1. {res_A['name']:<32} : 耗时 {res_A['latency']:.4f}s | 生成 {res_A['num_tokens']} Tokens | 吞吐率 {res_A['tps']:.2f} Tokens/s"
)
print(
    f"2. {res_B['name']:<32} : 耗时 {res_B['latency']:.4f}s | 生成 {res_B['num_tokens']} Tokens | 吞吐率 {res_B['tps']:.2f} Tokens/s"
)

tps_diff = ((res_B["tps"] - res_A["tps"]) / res_A["tps"]) * 100

print("-" * 60)
print(f"🚀 每秒生成吞吐率 (TPS) 提升幅度 : {tps_diff:+.2f}%")
print("=" * 60)