import json
from collections import defaultdict

prefix = "head_Llama-3.1-8B-Instruct_head_"
score = "prc"
target = "inst"
topks = "0 3.125p 6.25p 9.375p 12.5p 15.625p 18.75p 21.875p 25p 28.125p 31.25p 50p 62.5p 75p 100p".split()

attacks = ["none", "naive", "ignore", "escape_separation"]

# [metric][attack] -> list aligned with topks
attack_asr = defaultdict(list)
attack_vr  = defaultdict(list)

for t in topks:
    with open(f"{prefix}{score}_{target}_{t}.json", "r") as f:
        data = json.load(f)

    # 每個檔案內的 summary 可能順序不固定，先做成 dict 方便取
    by_attack = {s["attack"]: s for s in data["summary"]}

    for a in attacks:
        s = by_attack[a]
        attack_asr[a].append(s["attack_success_rate"])
        attack_vr[a].append(s["valid_rate"])

def fmt_topk(t):
    return t[:-1] if t.endswith("p") else t

def fmt_num(x):
    # 你範例是小數點後 9 位左右；這裡固定 9 位再去掉尾端 0
    s = f"{x:.9f}".rstrip("0").rstrip(".")
    return s

def print_table(title, metric_dict):
    print(title)
    print("\t".join(["topk"] + attacks))
    for i, t in enumerate(topks):
        row = [fmt_topk(t)]
        for a in attacks:
            row.append(fmt_num(metric_dict[a][i]))
        print("\t".join(row))

print_table("ASR", attack_asr)
print()
print_table("Valid", attack_vr)
