#!/usr/bin/env python3
"""Compare MiniMax H3 AdaLN modulation *outputs*, not curve-basis weights.

No torch/safetensors imports; reads .safetensors arrays via NumPy memmap.
This avoids PyTorch native crashes and only materializes one AdaLN layer
at a time. Does not modify source files.
"""

import argparse
import csv
import json
import math
import os
import struct
import numpy as np

PREFIXES = ("model.diffusion_model.", "diffusion_model.")
DTYPES = {"F16": np.dtype("<f2"), "BF16": np.dtype("<u2"), "F32": np.dtype("<f4")}


def normalize_key(k):
    for p in PREFIXES:
        if k.startswith(p):
            return k[len(p) :]
    return k


class TensorFile:
    def __init__(self, path):
        self.path = path
        with open(path, "rb") as f:
            hraw = f.read(8)
            if len(hraw) != 8:
                raise ValueError("Not a Safetensors file: " + path)
            length = struct.unpack("<Q", hraw)[0]
            if length > 500_000_000:
                raise ValueError("Unreasonably large Safetensors header")
            raw = f.read(length)
            if len(raw) != length:
                raise ValueError("Truncated Safetensors header")
        header = json.loads(raw)
        self.offset = 8 + length
        self.info = {}
        for key, info in header.items():
            if key == "__metadata__":
                continue
            nk = normalize_key(key)
            if nk in self.info:
                raise ValueError(f"Duplicate normalized key in {path}: {nk}")
            self.info[nk] = info

    def get(self, key):
        info = self.info[key]
        code = info["dtype"]
        if code not in DTYPES:
            raise ValueError(f"Unsupported tensor dtype {code}: {key}")
        dtype = DTYPES[code]
        shape = tuple(info["shape"])
        a, b = info["data_offsets"]
        expected_bytes = math.prod(shape) * dtype.itemsize
        if b - a != expected_bytes:
            raise ValueError(f"Invalid byte size for {key}")
        mmap = np.memmap(
            self.path, dtype=dtype, mode="r", offset=self.offset + a, shape=shape
        )
        if code == "BF16":
            result = (np.asarray(mmap, dtype=np.uint32) << 16).view(np.float32)
        else:
            result = np.asarray(mmap, dtype=np.float32)
        return result


def metrics(base, target):
    if base.shape != target.shape:
        raise ValueError(f"Shape mismatch: {base.shape} and {target.shape}")
    a = base.astype(np.float64, copy=False).ravel()
    b = target.astype(np.float64, copy=False).ravel()
    delta = b - a
    d2 = float(np.dot(delta, delta))
    a2 = float(np.dot(a, a))
    b2 = float(np.dot(b, b))
    cross = float(np.dot(a, b))
    rel = math.sqrt(d2 / a2) if a2 > 0 else (0.0 if d2 == 0 else float("inf"))
    denom = math.sqrt(a2 * b2)
    cos = max(-1.0, min(1.0, cross / denom)) if denom else (1.0 if d2 == 0 else 0.0)
    return rel * 100.0, cos, float(np.max(np.abs(delta)))


def evaluate(base_file, target_file, num_samples=9, num_blocks=50, include_final=True):
    bt = base_file.get("adaln_t_table")
    tt = target_file.get("adaln_t_table")
    if bt.shape != tt.shape or bt.ndim != 2:
        raise ValueError(f"Incompatible time tables: {bt.shape} / {tt.shape}")
    indices = np.unique(np.linspace(0, bt.shape[0] - 1, num_samples, dtype=np.int64))
    bv, tv = bt[indices], tt[indices]
    print(f"Timestep table shape: {bt.shape}", flush=True)
    print(f"Comparing rows: {list(map(int, indices))}", flush=True)
    table_rel, table_cos, table_max = metrics(bt, tt)
    print(
        f"RAW time table difference: {table_rel:.4f}% (cos {table_cos:.6f})", flush=True
    )
    print()

    names = [f"blocks.{i}.adaln_proj.linear" for i in range(num_blocks)]
    if include_final:
        names.append("final_layer.adaln_proj.linear")
    rows = []
    for i, stem in enumerate(names, 1):
        bw = base_file.get(stem + ".weight")
        tw = target_file.get(stem + ".weight")
        bb = base_file.get(stem + ".bias")
        tb = target_file.get(stem + ".bias")
        if bw.shape != tw.shape or bb.shape != tb.shape:
            raise ValueError(f"Unexpected AdaLN shape mismatch: {stem}")
        if bw.ndim != 2 or bw.shape[1] != bt.shape[1] or bb.shape != (bw.shape[0],):
            raise ValueError(f"Invalid AdaLN projection dimensions: {stem}")
        # Every model uses its OWN curve and OWN projection: these are the
        # quantities to compare, not its arbitrary 8D SVD coordinates.
        base_out = bv @ bw.T + bb
        target_out = tv @ tw.T + tb
        rel, cos, max_abs = metrics(base_out, target_out)
        wrel, wcos, _ = metrics(bw, tw)
        row = {
            "key": stem,
            "modulation_relerr_percent": rel,
            "modulation_cosine": cos,
            "modulation_max_abs": max_abs,
            "raw_weight_relerr_percent": wrel,
            "raw_weight_cosine": wcos,
        }
        rows.append(row)
        print(
            f"[{i:2}/{len(names)}] output rel={rel:9.4f}% cos={cos:.8f} | "
            f"raw W rel={wrel:9.3f}% | {stem}",
            flush=True,
        )
    return rows


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("base", help="Original Ref2VA checkpoint")
    ap.add_argument("target", help="FastH3 Omni checkpoint")
    ap.add_argument("--samples", type=int, default=9)
    ap.add_argument("--blocks", type=int, default=50)
    ap.add_argument("--no-final", action="store_true")
    ap.add_argument("--report", default="h3_omni_adaln_function.tsv")
    args = ap.parse_args()
    if args.samples < 2 or args.blocks < 0:
        ap.error("--samples must be >= 2 and --blocks must be >= 0")
    print("Using NumPy, no torch/CUDA. No checkpoint changes.\n", flush=True)
    rows = evaluate(
        TensorFile(args.base),
        TensorFile(args.target),
        args.samples,
        args.blocks,
        not args.no_final,
    )
    with open(args.report, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f, fieldnames=list(rows[0]) if rows else ["key"], delimiter="\t"
        )
        writer.writeheader()
        writer.writerows(rows)
    print("\nSUMMARY", flush=True)
    if rows:
        ordered = sorted(
            rows, key=lambda x: x["modulation_relerr_percent"], reverse=True
        )
        print(f"Compared layers: {len(rows)}")
        print(
            f"Mean output relative error: {sum(x['modulation_relerr_percent'] for x in rows)/len(rows):.5f}%"
        )
        print("Worst 10 output mismatches:")
        for x in ordered[:10]:
            print(
                f"  {x['modulation_relerr_percent']:9.5f}%  cos {x['modulation_cosine']:.7f}  {x['key']}"
            )
    print(f"Full report: {args.report}", flush=True)


if __name__ == "__main__":
    main()