import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import csv
import numpy as np

# ------------------ Load model ------------------
model_name = "Qwen/Qwen2-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

# ------------------ Detect architecture --------

print(model)
# Compute norms
def compute_norms(A: torch.Tensor):
    print(A.shape)
    norms = {}
    norms["1"]      = torch.norm(A, p=1).item()
    norms["inf"]    = torch.norm(A, p=float("inf")).item()
    norms["2"]      = torch.linalg.svdvals(A).max().item()
    norms["F"]      = torch.norm(A, p="fro").item()
    norms["*"]      = torch.linalg.svdvals(A).sum().item()
    norms["max"]    = A.abs().max().item()
    norms["2,1"]    = torch.norm(A, dim=0, p=2).sum().item()
    norms["2,1,T"]  = torch.norm(A.t(), dim=0, p=2).sum().item()
    return norms

def dict_of_lists(norm_list):
    result = {k: [] for k in norm_list[0].keys()}
    for d in norm_list:
        for k, v in d.items():
            result[k].append(v)
    return result

# Store norms and ratios
all_norms_Q, all_norms_K, norm_ratios = [], [], []

for layer in model.model.layers:
    W_q = layer.self_attn.q_proj.weight.detach()
    W_k = layer.self_attn.k_proj.weight.detach()
    if W_k.shape[0] != W_q.shape[0]:
        repeat_factor = W_q.shape[0] // W_k.shape[0]
        W_k = W_k.repeat(repeat_factor, 1)
    norms_q = compute_norms(W_q)
    norms_k = compute_norms(W_k)
    
    all_norms_Q.append(norms_q)
    all_norms_K.append(norms_k)
    
    ratio_dict = {k: norms_q[k] / norms_k[k] for k in norms_q.keys()}
    norm_ratios.append(ratio_dict)

Q_dict = dict_of_lists(all_norms_Q)
K_dict = dict_of_lists(all_norms_K)
R_dict = dict_of_lists(norm_ratios)

sections = [
    (["1", "inf", "2", "F"], ["1", "inf", "2", "F"]),
    (["*", "max", "2,1", "2,1,T"], ["*", "max", "2,1", "2,1,T"])
]

with open("qwen2.5_norms.csv", "w", newline="") as f:
    writer = csv.writer(f)

    for idx, (main_norms, keys) in enumerate(sections):
        # Header 1
        header1 = ["Layer"]
        for norm in main_norms:
            header1.extend([norm, "", ""])
        writer.writerow(header1)

        # Header 2
        header2 = [""]
        for _ in main_norms:
            header2.extend(["Q", "K", "Q/K"])
        writer.writerow(header2)

        # Per-layer rows
        for i in range(len(all_norms_Q)):
            row = [i + 1]
            for key in keys:
                row.extend([
                    round(Q_dict[key][i], 4),
                    round(K_dict[key][i], 4),
                    round(R_dict[key][i], 4)
                ])
            writer.writerow(row)

        # Mean and Std
        mean_row = ["Mean and Std"]
        std_row  = [""]

        for key in keys:
            q_vals = np.array(Q_dict[key])
            k_vals = np.array(K_dict[key])
            r_vals = np.array(R_dict[key])
            
            mean_row.extend([
                round(q_vals.mean(), 4),
                round(k_vals.mean(), 4),
                round(r_vals.mean(), 4)
            ])
            std_row.extend([
                round(q_vals.std(), 4),
                round(k_vals.std(), 4),
                round(r_vals.std(), 4)
            ])

        writer.writerow(mean_row)
        writer.writerow(std_row)