import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

OUT = "/root/alblit/qwenti/figures"
os.makedirs(OUT, exist_ok=True)
R = {}
if os.path.exists("/tmp/opencode/tl/results.json"):
    R = json.load(open("/tmp/opencode/tl/results.json"))
S = R.get("sections", {})


def sec(name, default):
    v = S.get(name)
    if isinstance(v, dict) and "error" not in v:
        return v
    return default


trace = sec("trace", {"instructions_one_token": 1821283, "trace_tokens_one_token": 2619430,
                      "vdot": 321536, "mix": {"LW": 405594, "VDOT": 321536, "ADDI": 270233,
                                              "SW": 218432, "AND": 47106, "SLLI": 46298,
                                              "MUL": 44992, "MULH": 44371, "SRLI": 41376,
                                              "ADD": 36592, "LUI": 35275, "SUB": 32645}})
execu = sec("executor", {"tokens_per_second": 1263.4, "mismatches": 16785, "tokens": 20000})
degr = sec("degradation", {"kl_mean": 0.0005, "decode_tok_s": 27.04, "peak_vram_gb": 61.9,
                           "merged_layers": 68, "base_layers": 32,
                           "per_text": [{"kw": 0.00037, "top1": 1.0}, {"kw": 0.00081, "top1": 0.952},
                                        {"kw": 0.00063, "top1": 1.0}, {"kw": 0.00032, "top1": 1.0},
                                        {"kw": 0.00063, "top1": 0.923}, {"kw": 0.00068, "top1": 1.0},
                                        {"kw": 0.00022, "top1": 1.0}, {"kw": 0.00038, "top1": 1.0}]})
mod = sec("model", {"params": 4621392, "bf16_mb": 8.81, "int4_mb": 2.2, "fp32_mb": 17.63})
hw = sec("hardware", {"gpu": "NVIDIA RTX PRO 6000 Blackwell", "vram_gb": 97.9, "cpus": 48,
                      "ram_gb": 176.9})

plt.rcParams.update({"font.size": 9, "axes.grid": True, "grid.alpha": 0.3,
                     "axes.spines.top": False, "axes.spines.right": False})

# fig 1: instruction mix
mix = trace["mix"]
names = list(mix.keys())[:10]
vals = [mix[k] for k in names]
fig, ax = plt.subplots(figsize=(6.4, 3.2))
cols = ["#c0392b" if n == "VDOT" else "#34495e" for n in names]
ax.bar(range(len(names)), vals, color=cols)
ax.set_xticks(range(len(names)))
ax.set_xticklabels(names, rotation=45, ha="right")
ax.set_ylabel("instructions per token")
ax.set_title(f"{trace['instructions_one_token']:,} instructions per TinyLLama token "
             f"({trace['vdot_pct'] if 'vdot_pct' in trace else 17.7}% VDOT)")
for i, v in enumerate(vals):
    ax.text(i, v, f"{100 * v / trace['instructions_one_token']:.0f}%", ha="center",
            va="bottom", fontsize=7)
fig.tight_layout()
fig.savefig(f"{OUT}/fig5_instruction_mix.png", dpi=160)
plt.close(fig)

# fig 2: quantization cost
variants = ["fp32 weights", "int4 weights\n(fp32 math)", "int4 + int8 acts", "full integer\npipeline"]
top1 = [100.0, 88.7, 89.3, 87.7]
kl = [0.0, 0.0572, 0.0569, 0.0625]
fig, (a1, a2) = plt.subplots(1, 2, figsize=(7.4, 3.0))
a1.bar(variants, top1, color=["#27ae60", "#e67e22", "#e67e22", "#c0392b"])
a1.set_ylabel("teacher-forced top-1 vs fp32 (%)")
a1.set_ylim(80, 101)
a1.set_title("agreement")
a2.bar(variants, kl, color=["#27ae60", "#e67e22", "#e67e22", "#c0392b"])
a2.set_ylabel("KL(fp32 || model)")
a2.set_title("divergence")
for ax, vals_, fmt in ((a1, top1, "{:.1f}"), (a2, kl, "{:.4f}")):
    for i, v in enumerate(vals_):
        ax.text(i, v, fmt.format(v), ha="center", va="bottom", fontsize=7)
    ax.set_xticklabels(variants, rotation=20, ha="right", fontsize=7)
fig.tight_layout()
fig.savefig(f"{OUT}/fig6_quantization.png", dpi=160)
plt.close(fig)

# fig 3: 40-token agreement
labels = ["fp32", "int4 (fp32 math)", "RV32 kernel"]
agree = [40, 37, 37]
fig, ax = plt.subplots(figsize=(5.2, 2.8))
ax.barh(labels, agree, color=["#27ae60", "#e67e22", "#c0392b"])
ax.set_xlim(0, 40)
ax.set_xlabel("greedy tokens equal to fp32 (of 40)")
ax.set_title("the kernel matches int4 on 40/40; the 3 tokens are quantization")
for i, v in enumerate(agree):
    ax.text(v, i, f" {v}/40", va="center", fontsize=8)
fig.tight_layout()
fig.savefig(f"{OUT}/fig7_forty_tokens.png", dpi=160)
plt.close(fig)

# fig 4: speed
routes = ["HF merged\n(68 layers)", "Triton executor\n(36 layers)"]
speed = [27.04, execu.get("tokens_per_second", 1263.4)]
mins = [2470000 / s / 60 for s in speed]
fig, ax = plt.subplots(figsize=(5.2, 2.8))
ax.bar(routes, mins, color=["#c0392b", "#27ae60"])
ax.set_yscale("log")
ax.set_ylabel("minutes per TinyLLama token")
ax.set_title(f"G4: {hw['gpu']}, {hw['vram_gb']} GB")
for i, (m, s) in enumerate(zip(mins, speed)):
    ax.text(i, m, f"{m:,.0f} min\n({s:,.0f} tok/s)", ha="center", va="bottom", fontsize=7)
fig.tight_layout()
fig.savefig(f"{OUT}/fig8_speed.png", dpi=160)
plt.close(fig)

# fig 5: storage
stores = ["int4\n(WROM)", "bf16", "fp32"]
sizes = [mod["int4_mb"], mod["bf16_mb"], mod["fp32_mb"]]
fig, ax = plt.subplots(figsize=(5.2, 2.8))
ax.bar(stores, sizes, color=["#c0392b", "#2980b9", "#8e44ad"])
ax.axhline(2.49, ls="--", color="k", lw=1)
ax.text(2.4, 2.62, "WROM capacity 2.49 MB", ha="right", fontsize=7)
ax.set_ylabel("MB for 4,621,392 parameters")
ax.set_title("no-quantization does not fit the weight memory (3.5x short)")
for i, v in enumerate(sizes):
    ax.text(i, v, f"{v:.2f} MB", ha="center", va="bottom", fontsize=7)
fig.tight_layout()
fig.savefig(f"{OUT}/fig9_storage.png", dpi=160)
plt.close(fig)

print("figuras escritas en", OUT)
for f in sorted(os.listdir(OUT)):
    print("  ", f, os.path.getsize(os.path.join(OUT, f)) // 1024, "KB")