from __future__ import annotations

import argparse
import json
import math
import os

import numpy as np
import torch
from PIL import Image

from dit import DiT

Image.MAX_IMAGE_PIXELS = None

def build(cfg):
    return DiT(dim=cfg["dim"], depth=cfg["depth"], heads=cfg["heads"])

def encode(ckpt, png_path, cfg_path, dim=384, depth=12, heads=6, which="ema"):
    ck = torch.load(ckpt, map_location="cpu")
    cfg = {"dim": dim, "depth": depth, "heads": heads}
    model = build(cfg)
    model.load_state_dict(ck[which])
    parts, manifest = [], []
    for name, p in model.named_parameters():
        a = p.detach().to(torch.float16).contiguous().view(-1).numpy()
        parts.append(a)
        manifest.append({"name": name, "shape": list(p.shape), "numel": int(a.size)})
    flat = np.concatenate(parts)
    N = flat.size
    side = math.ceil(math.sqrt(N))
    u16 = flat.view(np.uint16)
    img = np.zeros((side * side, 3), dtype=np.uint8)
    img[:N, 0] = (u16 >> 8).astype(np.uint8)
    img[:N, 1] = (u16 & 0xFF).astype(np.uint8)
    Image.fromarray(img.reshape(side, side, 3), "RGB").save(png_path)
    json.dump({"cfg": cfg, "params": manifest, "total_parameters": int(N),
               "side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused",
               "step": ck.get("step"), "val_loss": ck.get("val")},
              open(cfg_path, "w"))
    print(f"[png] encoded {N:,} params into a {side}x{side} PNG "
          f"({os.path.getsize(png_path)/1e6:.1f} MB)")
    return N, side

def load_model_png(png_path, cfg_path, device="cpu"):
    meta = json.load(open(cfg_path))
    model = build(meta["cfg"])
    arr = np.asarray(Image.open(png_path).convert("RGB")).reshape(-1, 3)
    total = meta["total_parameters"]
    hi = arr[:total, 0].astype(np.uint16)
    lo = arr[:total, 1].astype(np.uint16)
    flat = ((hi << 8) | lo).astype(np.uint16).view(np.float16)
    sd = dict(model.named_parameters())
    off = 0
    with torch.no_grad():
        for m in meta["params"]:
            n = m["numel"]
            chunk = flat[off:off + n].astype(np.float16)
            sd[m["name"]].copy_(torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32))
            off += n
    return model.to(device).eval()

def verify(ckpt, png_path, cfg_path, which="ema"):
    model = load_model_png(png_path, cfg_path)
    ref = torch.load(ckpt, map_location="cpu")[which]
    worst, name = 0.0, ""
    for k, p in model.named_parameters():
        d = (p.detach() - ref[k].float()).abs().max().item()
        if d > worst:
            worst, name = d, k
    scale = max(1e-12, ref[name].float().abs().max().item())
    print(f"[png] round-trip max abs error {worst:.3e} on {name} (relative {worst/scale:.2e})")
    return worst

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="/root/runs/pm5/best.pt")
    ap.add_argument("--png", default="model.png")
    ap.add_argument("--manifest", default="model_png.json")
    args = ap.parse_args()
    encode(args.ckpt, args.png, args.manifest)
    verify(args.ckpt, args.png, args.manifest)

if __name__ == "__main__":
    main()