#!/usr/bin/env python3
"""Compare primary-image IEEE float values in two simple FITS files.

The renderer writes a single 32-bit RGB primary HDU.  Keeping this dependency-
free tool in the repository makes HDR reference comparisons available on the
minimal Gentoo installation as well as in CI.
"""

import argparse
import math
import struct
import sys


def read_fits(path):
    with open(path, "rb") as stream:
        cards = []
        while True:
            block = stream.read(2880)
            if len(block) != 2880:
                raise ValueError(f"{path}: truncated FITS header")
            cards.extend(block[i:i + 80].decode("ascii") for i in range(0, 2880, 80))
            if any(card.startswith("END") for card in cards[-36:]):
                break
        values = {}
        for card in cards:
            if card[8:10] == "= ":
                values[card[:8].strip()] = card[10:30].strip()
        if values.get("BITPIX") != "-32" or values.get("NAXIS") != "3":
            raise ValueError(f"{path}: expected a 32-bit, three-axis primary image")
        shape = tuple(int(values[f"NAXIS{axis}"]) for axis in (1, 2, 3))
        count = math.prod(shape)
        payload = stream.read(count * 4)
        if len(payload) != count * 4:
            raise ValueError(f"{path}: truncated image payload")
        return shape, payload


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("reference")
    parser.add_argument("candidate")
    parser.add_argument("--abs-tolerance", type=float, default=0.0)
    parser.add_argument("--rel-tolerance", type=float, default=0.0)
    args = parser.parse_args()
    if args.abs_tolerance < 0 or args.rel_tolerance < 0:
        parser.error("tolerances must be non-negative")
    ref_shape, ref_payload = read_fits(args.reference)
    candidate_shape, candidate_payload = read_fits(args.candidate)
    if ref_shape != candidate_shape:
        print(f"shape mismatch: {ref_shape} != {candidate_shape}", file=sys.stderr)
        return 1
    maximum_abs = maximum_rel = 0.0
    mismatch_count = 0
    worst_index = 0
    for index, (reference, candidate) in enumerate(
            zip(struct.iter_unpack(">f", ref_payload),
                struct.iter_unpack(">f", candidate_payload))):
        reference, candidate = reference[0], candidate[0]
        if not math.isfinite(reference) or not math.isfinite(candidate):
            equal = reference == candidate
            absolute = relative = math.inf if not equal else 0.0
        else:
            absolute = abs(candidate - reference)
            relative = absolute / max(abs(reference), 1.0e-30)
            equal = absolute <= args.abs_tolerance or relative <= args.rel_tolerance
        if absolute > maximum_abs:
            maximum_abs, worst_index = absolute, index
        maximum_rel = max(maximum_rel, relative)
        mismatch_count += not equal
    print(f"floatdiff: shape={ref_shape} samples={math.prod(ref_shape)} "
          f"mismatches={mismatch_count} max_abs={maximum_abs:.9g} "
          f"max_rel={maximum_rel:.9g} worst_sample={worst_index}")
    return 1 if mismatch_count else 0


if __name__ == "__main__":
    sys.exit(main())