"""Revalidate final decoder arithmetic using saved, unchanged DiT latents.

Retains original outputs in an ignored development directory. Separate decode
measurements are never presented as a freshly measured end-to-end request.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import time
from pathlib import Path

import mlx.core as mx

from h3_mlx.runtime import configure_memory
from h3_mlx.vae import VAE, to_image


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--vae", required=True)
    p.add_argument("--runs", default="outputs/suite")
    a = p.parse_args()
    configure_memory()
    root = Path(a.runs)
    model = VAE(a.vae)
    for trace in sorted(root.rglob("*.trace.safetensors")):
        image = trace.with_name(trace.name.replace(".trace.safetensors", ".png"))
        record = image.with_suffix(".json")
        if not record.exists():
            continue
        row = json.loads(record.read_text())
        if row.get("decoder_revalidated"):
            continue
        backup = Path("outputs/pre-layernorm-fix") / image.relative_to(root)
        backup.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(image, backup)
        shutil.copyfile(record, backup.with_suffix(".json"))
        arrays = mx.load(str(trace))
        latent = arrays["latent"]
        mx.eval(latent)
        del arrays
        model.tile_batch = row["vae"]["tile_batch"]
        start = time.perf_counter()
        pixels = model.decode(latent)
        result = to_image(pixels, *row["size"])
        elapsed = time.perf_counter() - start
        result.save(image)
        original = row["timings"].copy()
        row["source_request_timings_before_decoder_fix"] = original
        row["timings"].pop("total_s", None)
        row["timings"]["decode_s"] = elapsed
        row["decoder_revalidated"] = {
            "revision": 2,
            "method": "FP32 LayerNorm affine before BF16 rounding",
            "same_saved_latent": True,
            "decode_s": elapsed,
        }
        row["timing_scope"] = (
            "Saved DiT stages plus separately timed corrected decoder; no fresh end-to-end total"
        )
        row["png_sha256"] = hashlib.sha256(image.read_bytes()).hexdigest()
        record.write_text(json.dumps(row, indent=2) + "\n")
        print(json.dumps({"image": str(image), "decode_s": elapsed}), flush=True)
    model.clear()


if __name__ == "__main__":
    main()