import torch
import numpy as np
import csv
from transformers import ViTModel

# Load ViT
model = ViTModel.from_pretrained("google/vit-base-patch16-224")

hidden_size = model.config.hidden_size
num_heads = model.config.num_attention_heads
head_dim = hidden_size // num_heads

# Norm calculator
def compute_norms(A: torch.Tensor):
    return {
        "1": torch.norm(A, p=1).item(),
        "F": torch.norm(A, p="fro").item(),
        "*": torch.linalg.svdvals(A).sum().item(),
        "2,1": torch.norm(A, dim=0, p=2).sum().item(),
        "2,1,T": torch.norm(A.t(), dim=0, p=2).sum().item()
    }

norm_names = ["1","F","*","2,1","2,1,T"]
subcols = ["Q","K","Q/K","V","O","V/O"]

outfile = "vit_qkvo_norms.csv"

with open(outfile, "w", newline="") as f:
    writer = csv.writer(f)

    for layer_idx, layer in enumerate(model.encoder.layer, start=1):
        # Header rows
        header1 = [f"Layer {layer_idx}"]
        for n in norm_names:
            header1.extend([n,"","","","",""])
        writer.writerow(header1)

        header2 = [""]
        for _ in norm_names:
            header2.extend(subcols)
        writer.writerow(header2)

        rows = []

        # Extract weights + biases
        W_q = torch.cat([layer.attention.attention.query.weight.detach(),
                         layer.attention.attention.query.bias.detach().unsqueeze(1)], dim=1)
        W_k = torch.cat([layer.attention.attention.key.weight.detach(),
                         layer.attention.attention.key.bias.detach().unsqueeze(1)], dim=1)
        W_v = torch.cat([layer.attention.attention.value.weight.detach(),
                         layer.attention.attention.value.bias.detach().unsqueeze(1)], dim=1)
        W_o = torch.cat([layer.attention.output.dense.weight.detach(),
                         layer.attention.output.dense.bias.detach().unsqueeze(1)], dim=1)

        # Reshape into heads
        W_q_heads = W_q.view(num_heads, head_dim, -1)
        W_k_heads = W_k.view(num_heads, head_dim, -1)
        W_v_heads = W_v.view(num_heads, head_dim, -1)
        W_o_heads = W_o.view(num_heads, head_dim, -1)

        for h in range(num_heads):
            row = [f"Head {h+1}"]
            for norm in norm_names:
                nq = compute_norms(W_q_heads[h])[norm]
                nk = compute_norms(W_k_heads[h])[norm]
                nv = compute_norms(W_v_heads[h])[norm]
                no = compute_norms(W_o_heads[h])[norm]

                qk_ratio = nq/(nk+1e-12)
                vo_ratio = nv/(no+1e-12)

                row.extend([
                    round(nq,4),
                    round(nk,4),
                    round(qk_ratio,4),
                    round(nv,4),
                    round(no,4),
                    round(vo_ratio,4)
                ])
            writer.writerow(row)
            rows.append(row[1:])

        # Mean & Std rows
        arr = np.array(rows, dtype=float)
        mean = np.round(arr.mean(axis=0),4)
        std = np.round(arr.std(axis=0),4)

        writer.writerow(["Mean"] + mean.tolist())
        writer.writerow(["Std"] + std.tolist())
        writer.writerow([])

print(f"✅ Saved CSV: {outfile}")