#!/usr/bin/env python3
"""Compare two MiniMax H3 Safetensors checkpoints without torch or safetensors bindings.

Only Python's JSON/file I/O and NumPy are used. One small block from each
checkpoint is read at a time. This avoids torch native crashes and holding
large transformer weights in RAM. Original checkpoints are never modified.
"""

import argparse
import json
import math
import os
import struct
from collections import Counter
from pathlib import Path

import numpy as np

PREFIXES = ("model.diffusion_model.", "diffusion_model.")

# Safetensors stores everything in little-endian byte order.
DTYPES = {
    "F64": ("<f8", 8),
    "F32": ("<f4", 4),
    "F16": ("<f2", 2),
    "BF16": ("<u2", 2),
    "I8": ("i1", 1),
    "U8": ("u1", 1),
    "BOOL": ("?", 1),
    "I16": ("<i2", 2),
    "U16": ("<u2", 2),
    "I32": ("<i4", 4),
    "U32": ("<u4", 4),
    "I64": ("<i8", 8),
    "U64": ("<u8", 8),
}


def normalize_key(key):
    for prefix in PREFIXES:
        if key.startswith(prefix):
            return key[len(prefix) :]
    return key


class Checkpoint:
    def __init__(self, path):
        self.path = Path(path)
        self.file = self.path.open("rb")
        try:
            raw_len = self.file.read(8)
            if len(raw_len) != 8:
                raise ValueError(f"Invalid Safetensors header in {path}")
            hlen = struct.unpack("<Q", raw_len)[0]
            if hlen > 512 * 1024 * 1024:
                raise ValueError("Header unexpectedly large")
            raw_header = self.file.read(hlen)
            if len(raw_header) != hlen:
                raise ValueError(f"Truncated Safetensors header in {path}")
            header = json.loads(raw_header)
            self.data_start = 8 + hlen
            self.entries = {}
            max_bytes = self.path.stat().st_size - self.data_start
            for raw_key, info in header.items():
                if raw_key == "__metadata__":
                    continue
                key = normalize_key(raw_key)
                if key in self.entries:
                    raise ValueError(f"Duplicate normalized tensor key: {key}")
                dtype = info["dtype"]
                if dtype not in DTYPES:
                    raise ValueError(f"Unsupported dtype {dtype} for {key}")
                shape = tuple(info["shape"])
                start, end = info["data_offsets"]
                count = math.prod(shape)
                nbytes = count * DTYPES[dtype][1]
                if start < 0 or end < start or end > max_bytes or end - start != nbytes:
                    raise ValueError(f"Invalid offsets or shape for {key}")
                self.entries[key] = {
                    "raw_key": raw_key,
                    "dtype": dtype,
                    "shape": shape,
                    "start": start,
                    "count": count,
                }
        except BaseException:
            self.file.close()
            raise

    def close(self):
        self.file.close()

    def read_chunk(self, key, element_offset, n):
        entry = self.entries[key]
        dtype = entry["dtype"]
        width = DTYPES[dtype][1]
        offset = self.data_start + entry["start"] + element_offset * width
        self.file.seek(offset)
        data = self.file.read(n * width)
        if len(data) != n * width:
            raise IOError(f"Unexpected EOF reading {key}")
        return data

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


def to_float32(raw, dtype):
    values = np.frombuffer(raw, dtype=np.dtype(DTYPES[dtype][0]))
    if dtype == "BF16":
        # BF16 consists of the top 16 bits of an IEEE754 FP32 value.
        return np.left_shift(values.astype(np.uint32), 16).view(np.float32)
    return values.astype(np.float32, copy=False)


def cast_f32_to_dtype(values, dtype):
    """Float32 representation after rounding values to a storage dtype."""
    if dtype == "F32":
        return values
    if dtype == "F16":
        return values.astype(np.float16).astype(np.float32)
    if dtype == "BF16":
        # Round-to-nearest-even BF16 encoding for finite FP32 values.
        bits = values.astype(np.float32, copy=False).view(np.uint32)
        rounded = bits + (np.uint32(0x7FFF) + ((bits >> 16) & 1))
        return (rounded & np.uint32(0xFFFF0000)).view(np.float32)
    if dtype == "F64":
        return values  # comparison operates in FP32
    np_dtype = np.dtype(DTYPES[dtype][0])
    return values.astype(np_dtype).astype(np.float32)


def array_equal(a, b):
    return np.array_equal(a, b, equal_nan=True)


def compare_tensor(base, target, key, chunk_elements):
    be = base.entries[key]
    te = target.entries[key]
    if be["shape"] != te["shape"]:
        raise ValueError(f"Shape mismatch: {key}")

    identical_bytes = be["dtype"] == te["dtype"]
    base_casts_to_target = True
    target_casts_to_base = True
    diff_sq = base_sq = target_sq = dot = 0.0
    max_abs = 0.0
    processed = 0
    for start in range(0, be["count"], chunk_elements):
        n = min(chunk_elements, be["count"] - start)
        braw = base.read_chunk(key, start, n)
        traw = target.read_chunk(key, start, n)

        if identical_bytes and braw != traw:
            identical_bytes = False
        b = to_float32(braw, be["dtype"])
        t = to_float32(traw, te["dtype"])

        if not np.isfinite(b).all() or not np.isfinite(t).all():
            raise ValueError(f"Non-finite values in tensor {key} at offset {start}")
        if base_casts_to_target and not array_equal(
            cast_f32_to_dtype(b, te["dtype"]), t
        ):
            base_casts_to_target = False
        if target_casts_to_base and not array_equal(
            cast_f32_to_dtype(t, be["dtype"]), b
        ):
            target_casts_to_base = False

        diff = t - b
        diff_sq += float(np.einsum("i,i->", diff, diff, dtype=np.float64))
        base_sq += float(np.einsum("i,i->", b, b, dtype=np.float64))
        target_sq += float(np.einsum("i,i->", t, t, dtype=np.float64))
        dot += float(np.einsum("i,i->", b, t, dtype=np.float64))
        max_abs = max(max_abs, float(np.max(np.abs(diff))))
        processed += n

    dnorm = math.sqrt(diff_sq)
    bnorm = math.sqrt(base_sq)
    relerr = dnorm / bnorm if bnorm else (0.0 if not dnorm else float("inf"))
    denom = math.sqrt(base_sq * target_sq)
    cosine = dot / denom if denom else (1.0 if not diff_sq else 0.0)
    cosine = min(1.0, max(-1.0, cosine))
    if identical_bytes:
        classification = "IDENTICAL"
    elif be["dtype"] != te["dtype"] and (base_casts_to_target or target_casts_to_base):
        classification = "CAST_EQ"
    else:
        classification = "CHANGED"
    return {
        "key": key,
        "shape": be["shape"],
        "base_dtype": be["dtype"],
        "target_dtype": te["dtype"],
        "classification": classification,
        "relerr": relerr,
        "cosine": cosine,
        "max_abs": max_abs,
        "elements": processed,
    }


def main():
    parser = argparse.ArgumentParser(
        description="Compare MiniMax H3 checkpoint values without PyTorch or CUDA"
    )
    parser.add_argument("base", help="Base Ref2VA safetensors checkpoint")
    parser.add_argument("target", help="FastH3 Omni safetensors checkpoint")
    parser.add_argument("--report", default="h3_omni_value_comparison.tsv")
    parser.add_argument("--chunk-elements", type=int, default=1_000_000)
    args = parser.parse_args()
    if args.chunk_elements < 1:
        parser.error("--chunk-elements must be >= 1")

    with Checkpoint(args.base) as base, Checkpoint(args.target) as target:
        shared = sorted(set(base.entries) & set(target.entries))
        print(f"Shared tensors: {len(shared)}", flush=True)
        print("Using NumPy + direct Safetensors reads (no torch)\n", flush=True)
        results = []
        skipped = []
        for i, key in enumerate(shared, 1):
            bshape = base.entries[key]["shape"]
            tshape = target.entries[key]["shape"]
            if bshape != tshape:
                skipped.append(key)
                print(
                    f"[{i:3}/{len(shared)}] SHAPE     {key}: {bshape} -> {tshape}",
                    flush=True,
                )
                continue
            print(f"[{i:3}/{len(shared)}] START     {key}", flush=True)
            result = compare_tensor(base, target, key, args.chunk_elements)
            results.append(result)
            print(
                f"[{i:3}/{len(shared)}] {result['classification']:9} "
                f"rel={result['relerr']*100:9.5f}% "
                f"cos={result['cosine']:.8f} {key}",
                flush=True,
            )

    counter = Counter(r["classification"] for r in results)
    changed = sorted(
        (r for r in results if r["classification"] == "CHANGED"),
        key=lambda r: r["relerr"],
        reverse=True,
    )
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print(f"Comparable tensors: {len(results)}")
    print(f"Shape mismatches:   {len(skipped)}")
    print(f"Identical:          {counter['IDENTICAL']}")
    print(f"Dtype/cast only:    {counter['CAST_EQ']}")
    print(f"Actually changed:   {counter['CHANGED']}")
    print("\nTOP 30 CHANGED")
    for item in changed[:30]:
        print(f"{item['relerr']*100:10.6f}% cos {item['cosine']:.8f} {item['key']}")

    report = Path(args.report)
    with report.open("w", encoding="utf-8", newline="") as output:
        output.write(
            "classification\trelerr_percent\tcosine\tmax_abs\tbase_dtype\ttarget_dtype\tshape\tkey\n"
        )
        for item in sorted(results, key=lambda x: x["relerr"], reverse=True):
            output.write(
                f"{item['classification']}\t{item['relerr']*100:.10f}\t"
                f"{item['cosine']:.10f}\t{item['max_abs']:.10e}\t"
                f"{item['base_dtype']}\t{item['target_dtype']}\t"
                f"{item['shape']}\t{item['key']}\n"
            )
    print(f"\nFull report: {report.resolve()}")
    print(
        "Note: CAST_EQ means identical after a dtype cast, not proof of no training change."
    )


if __name__ == "__main__":
    main()