#!/usr/bin/env python3

######################################
# author: generated by Claude Sonnet 5
######################################

"""
hipo2npz-diff - Compare two NumPy .npz files and report differences.

Usage:
    hipo2npz-diff file1.npz file2.npz
    hipo2npz-diff file1.npz file2.npz --rtol 1e-5 --atol 1e-8
    hipo2npz-diff file1.npz file2.npz --exact
    hipo2npz-diff file1.npz file2.npz --verbose

Exit codes:
    0 - files are equivalent (given tolerance)
    1 - differences found
    2 - error (bad file, etc.)
"""

import argparse
import sys

import numpy as np


def diff_npz(path_a, path_b, rtol, atol, exact, verbose):
    try:
        a = np.load(path_a, allow_pickle=True)
        b = np.load(path_b, allow_pickle=True)
    except Exception as e:
        print(f"Error loading files: {e}", file=sys.stderr)
        sys.exit(2)

    keys_a = set(a.files)
    keys_b = set(b.files)

    only_in_a = sorted(keys_a - keys_b)
    only_in_b = sorted(keys_b - keys_a)
    common = sorted(keys_a & keys_b)

    has_diff = False

    if only_in_a:
        has_diff = True
        print(f"Keys only in {path_a}:")
        for k in only_in_a:
            print(f"  - {k}")

    if only_in_b:
        has_diff = True
        print(f"Keys only in {path_b}:")
        for k in only_in_b:
            print(f"  - {k}")

    for key in common:
        arr_a, arr_b = a[key], b[key]

        if arr_a.shape != arr_b.shape:
            has_diff = True
            print(f"[{key}] shape mismatch: {arr_a.shape} vs {arr_b.shape}")
            continue

        if arr_a.dtype != arr_b.dtype and verbose:
            print(f"[{key}] dtype differs: {arr_a.dtype} vs {arr_b.dtype}")

        is_float = np.issubdtype(arr_a.dtype, np.floating)

        try:
            if exact:
                # equal_nan=True: NaNs in the same position count as matching,
                # since NaN != NaN by IEEE rules but that's rarely what you want here
                equal = np.array_equal(arr_a, arr_b, equal_nan=is_float)
            else:
                equal = np.allclose(arr_a, arr_b, rtol=rtol, atol=atol, equal_nan=True)
        except TypeError:
            # Non-numeric / object arrays: fall back to plain equality (no equal_nan support)
            equal = np.array_equal(arr_a, arr_b)

        if not equal:
            has_diff = True
            print(f"[{key}] values differ", end="")
            try:
                fa = arr_a.astype(np.float64)
                fb = arr_b.astype(np.float64)
                diff = np.abs(fa - fb)

                # Positions where exactly one side is NaN (a "real" mismatch, not just
                # matching NaNs) vs. positions with a finite numeric difference
                nan_mismatch = np.isnan(fa) != np.isnan(fb)
                finite_mask = ~np.isnan(fa) & ~np.isnan(fb)

                n_diff = np.count_nonzero(
                    finite_mask & (diff > (atol + rtol * np.abs(fb)))
                )
                n_nan_mismatch = np.count_nonzero(nan_mismatch)

                max_diff = np.max(diff[finite_mask]) if finite_mask.any() else 0.0
                print(
                    f" (max abs diff = {max_diff:.6g}, "
                    f"{n_diff}/{arr_a.size} elements differ, "
                    f"{n_nan_mismatch} NaN-mismatch positions)"
                )
            except (TypeError, ValueError):
                print()
        elif verbose:
            print(f"[{key}] OK (identical within tolerance)")

    if not has_diff:
        print(f"No differences found between {path_a} and {path_b}"
              + ("" if exact else f" (rtol={rtol}, atol={atol})"))

    return has_diff


def main():
    parser = argparse.ArgumentParser(description="Diff two .npz files.")
    parser.add_argument("file_a", help="First .npz file")
    parser.add_argument("file_b", help="Second .npz file")
    parser.add_argument("--rtol", type=float, default=1e-5, help="Relative tolerance for float comparison (default: 1e-5)")
    parser.add_argument("--atol", type=float, default=1e-8, help="Absolute tolerance for float comparison (default: 1e-8)")
    parser.add_argument("--exact", action="store_true", help="Require exact equality instead of tolerance-based comparison")
    parser.add_argument("--verbose", action="store_true", help="Print status for matching keys too")
    args = parser.parse_args()

    has_diff = diff_npz(args.file_a, args.file_b, args.rtol, args.atol, args.exact, args.verbose)
    sys.exit(1 if has_diff else 0)


if __name__ == "__main__":
    main()
