from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig
from alphalora.expert_number import calculate_expert
import json
import torch
from lsaq_quant_assign import quantize_llama_like
import os 
import argparse
import numpy as np

parser = argparse.ArgumentParser(description="parser")
parser.add_argument("--metric_name", type=str, default="alpha")
parser.add_argument("--model_id", type=str, default="Llama-2-7b-hf")
parser.add_argument("--keyword", type=str, default=None)
parser.add_argument("--cuda_id", type=int, default=1)

args = parser.parse_args()

# model_names = ['Qwen3-8B', ''Llama-2-7b-hf'', 'Mistral-7B-Instruct-v0.3','Llama-3.2-3B-Instruct']
model_ids = [args.model_id]


for model_id in model_ids:
    # Llama-3.2-3B-Instruct
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map=f"cuda:{args.cuda_id}",
        torch_dtype=torch.float16
    )
    model_name = os.path.basename(model_id.rstrip("/"))
    tokenizer = AutoTokenizer.from_pretrained(model_id)

    prompt = "Once upon a time"
    inputs = tokenizer(prompt, return_tensors="pt").to(f"cuda:{args.cuda_id}")

    # 生成文本
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=True,        # 随机采样（非贪婪）
            top_k=50,              # 限制采样候选
            top_p=0.95,            # nucleus sampling
            temperature=0.7        # 控制生成多样性
        )

    # 解码为字符串
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(decoded)
    print(model)
    # 计算 alpha_values

    # 'alpha': fix_finger, # alpha_values
    # 'alpha_hat': fix_finger_hat,
    # 'stable_rank': stable_rank,
    # 'effective_rank': effective_rank,
    # 'ZD': ZD,

    metric_name = args.metric_name
    keyword=args.keyword
    all_layer_alpha = calculate_expert(model, metric=metric_name, keyword=keyword)
    
    value_index_pairs = [(value, idx) for idx, value in enumerate(all_layer_alpha)]
    sorted_pairs = sorted(value_index_pairs, key=lambda x: x[0], reverse=True)
    print(sorted_pairs)
    print(f"metric_name {metric_name}:", [pairs[1] for pairs in sorted_pairs])

    # scores = np.array(all_layer_alpha)
    # normalized_scores = (scores - scores.min()) / (scores.max() - scores.min())
    # print(f"normalized_scores: {normalized_scores}")

    os.makedirs(f"metrics/{model_name}", exist_ok=True)
    with open(f"metrics/{model_name}/{metric_name}.json", "w") as f:
        json.dump(all_layer_alpha, f, indent=4)


# for bit in [8]:
#     print(bit)
#     layer_to_quant = list(range(32))
#     print(layer_to_quant)
#     mlp_quant = [f'layers.{item}.mlp' for item in layer_to_quant]
#     self_attn_quant = [f'layers.{item}.self_attn' for item in layer_to_quant]

#     print(f'quanting ... ')
#     model_lsaq = quantize_llama_like(model, mlp_quant, self_attn_quant, 4)
#     print(f'quanted')
#     print(model_lsaq)

#     with torch.no_grad():
#         outputs = model.generate(
#             **inputs,
#             max_new_tokens=100,
#             do_sample=True,        # 随机采样（非贪婪）
#             top_k=50,              # 限制采样候选
#             top_p=0.95,            # nucleus sampling
#             temperature=0.7        # 控制生成多样性
#         )

#     # 解码为字符串
#     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
#     print(decoded)

    

#     # save_dir = f"../models/Llama-2-7b-hf-qint{bit}"
#     # model_lsaq.save_pretrained(save_dir)
#     # tokenizer.save_pretrained(save_dir)