main.py

#!/usr/bin/env python3
"""
Compare two binary files with the following layout:

Header:
- 7 bytes: ASCII magic string "MWAOCAL"
- 10 bytes: zeros (b'\x00' * 10)
- uint32 LE: num_timeblocks
- uint32 LE: total_num_tiles
- uint32 LE: total_num_chanblocks
- uint32 LE: num_polarisations
- float64 LE: start_time (GPS seconds)
- float64 LE: end_time (GPS seconds)

Data:
- float64 array with shape:
    (num_timeblocks,
     total_num_tiles,
     total_num_chanblocks,
     2 * num_polarisations)

The script validates headers/metadata and compares the data stream in chunks.
"""

import argparse
from io import BufferedReader
import os
import struct
from typing import Any
import numpy as np

MAGIC = b"MWAOCAL\0"
ZERO_PAD_LEN = 8
UINT32_LE = "<I"
FLOAT64_LE = "<d"

# Header offsets and sizes (fixed)
HEADER_FIXED_LEN = len(MAGIC) + ZERO_PAD_LEN  # 17 bytes
# 4 * uint32 + 2 * float64
HEADER_META_FMT = "<I I I I d d"
HEADER_META_LEN = struct.calcsize(HEADER_META_FMT)
HEADER_TOTAL_LEN = HEADER_FIXED_LEN + HEADER_META_LEN  # 17 + 32 = 49 bytes


def read_header(fp):
    """
    Read and validate header; return a dict with fields.
    Raises ValueError on format violations.
    """
    # Read magic + zero pad
    fixed = fp.read(HEADER_FIXED_LEN)
    if len(fixed) != HEADER_FIXED_LEN:
        raise ValueError(f"File too small: expected >= {HEADER_TOTAL_LEN} bytes, got {len(fixed)} in fixed header.")

    magic = fixed[: len(MAGIC)]
    if magic != MAGIC:
        raise ValueError(f"Bad magic: expected {MAGIC!r}, got {magic!r}")

    zero_pad = fixed[len(MAGIC) :]
    if zero_pad != b"\x00" * ZERO_PAD_LEN:
        raise ValueError(f"Zero pad check failed: expected {ZERO_PAD_LEN} zero bytes. Got {zero_pad!r}")

    meta_bytes = fp.read(HEADER_META_LEN)
    if len(meta_bytes) != HEADER_META_LEN:
        raise ValueError("Incomplete metadata section.")

    (num_timeblocks, total_num_tiles, total_num_chanblocks, num_polarisations, start_time, end_time) = struct.unpack(
        HEADER_META_FMT, meta_bytes
    )

    # Basic sanity checks
    for name, val in [
        ("num_timeblocks", num_timeblocks),
        ("total_num_tiles", total_num_tiles),
        ("total_num_chanblocks", total_num_chanblocks),
        ("num_polarisations", num_polarisations),
    ]:
        if val == 0:
            raise ValueError(f"Invalid {name}: must be > 0, got {val}")

    return {
        "num_timeblocks": num_timeblocks,
        "total_num_tiles": total_num_tiles,
        "total_num_chanblocks": total_num_chanblocks,
        "num_polarisations": num_polarisations,
        "start_time": start_time,
        "end_time": end_time,
    }


def expected_data_bytes(meta):
    """
    Compute expected byte length of the data section based on metadata.
    """
    n_time = meta["num_timeblocks"]
    n_tiles = meta["total_num_tiles"]
    n_chan = meta["total_num_chanblocks"]
    n_pol = meta["num_polarisations"]
    n_last = 2 * n_pol  # as specified

    total_elems = n_time * n_tiles * n_chan * n_last
    return total_elems * 8  # float64 = 8 bytes


def compare_headers(meta1, meta2):
    """
    Compare metadata dictionaries. Returns list of differences.
    """
    diffs = []
    keys = ["num_timeblocks", "total_num_tiles", "total_num_chanblocks", "num_polarisations", "start_time", "end_time"]
    for k in keys:
        v1, v2 = meta1[k], meta2[k]
        if v1 != v2:
            diffs.append((k, v1, v2))
    return diffs


def compare_asymmetrical_data_stream(
    meta1: dict[str, Any],
    meta2: dict[str, Any],
    fp1: BufferedReader,
    fp2: BufferedReader,
    f1_bytes_to_compare,
    f2_bytes_to_compare,
):
    n_time1 = meta1["num_timeblocks"]
    n_tiles1 = meta1["total_num_tiles"]
    n_chan1 = meta1["total_num_chanblocks"]
    n_pol1 = meta1["num_polarisations"]

    n_time2 = meta2["num_timeblocks"]
    n_tiles2 = meta2["total_num_tiles"]
    n_chan2 = meta2["total_num_chanblocks"]
    n_pol2 = meta2["num_polarisations"]

    f1_elem_size = 8 * n_chan1  # f64 per fine chan
    f2_elem_size = 8 * n_chan2  # f64 per fine chan
    chan_ratio = n_chan1 // n_chan2 if n_chan1 > n_chan2 else n_chan2 // n_chan1
    print(f"Channel ratio: {chan_ratio}:1")

    # Move file pointers to the start of the data region
    fp1.seek(HEADER_TOTAL_LEN)
    fp2.seek(HEADER_TOTAL_LEN)

    # Read all data into memory (caution: may be large!)
    data1 = np.fromfile(fp1, dtype="<f8")
    data2 = np.fromfile(fp2, dtype="<f8")

    f1_elements = f1_bytes_to_compare // f1_elem_size
    f2_elements = f2_bytes_to_compare // f2_elem_size

    assert f1_elements == f2_elements, "Files don't have same number of elements!"

    arr1 = data1.reshape(n_time1, n_tiles1, n_chan1, 2 * n_pol1)
    arr2 = data2.reshape(n_time2, n_tiles2, n_chan2, 2 * n_pol2)

    print(f"  File 1 shape: {arr1.shape}")
    print(f"  File 2 shape: {arr2.shape}")

    # Find the larger chanblocks dimension
    if n_chan1 > n_chan2:
        reshaped = arr1.reshape(n_time1, n_tiles1, n_chan1 // chan_ratio, chan_ratio, n_pol1 * 2)
        arr1_new = reshaped.mean(axis=3)
        arr2_new = arr2
        print(arr1_new.shape, arr2_new.shape)
    else:
        reshaped = arr2.reshape(n_time1, n_tiles1, n_chan1 // chan_ratio, chan_ratio, n_pol1 * 2)
        arr2_new = reshaped.mean(axis=3)
        arr1_new = arr1
        print(arr1_new.shape, arr2_new.shape)

    matches = 0
    mismatches = 0

    output_type = 2

    if output_type == 1:
        # Compare value by value
        for t in range(n_time1):
            for tile in range(n_tiles1):
                for pol in range(n_pol1 * 2):
                    # Compare channels
                    for chan in range(min(n_chan1, n_chan2)):
                        val1 = arr1_new[t, tile, chan, pol]
                        val2 = arr2_new[t, tile, chan, pol]
                        if np.isnan(val1) and np.isnan(val2):
                            # This is ok
                            matches += 1
                            # print(f"Match at time {t}, tile {tile}, chan {chan}, pol {pol}: both NaN")
                            continue
                        if not np.isclose(val1, val2, rtol=1e-8, atol=1e-12):
                            print(
                                f"Mismatch at time {t}, tile {tile}, chan {chan}, pol {pol}: {arr1[t, tile, (chan * chan_ratio) : (chan * chan_ratio) + chan_ratio, pol]} => {val1} != {val2}"
                            )
                            mismatches += 1
                        else:
                            # print(f"Match at time {t}, tile {tile}, chan {chan}, pol {pol}: {val1} == {val2}")
                            matches += 1
    elif output_type == 2:
        # Compare time, tile, chan for all 8 pols
        for t in range(n_time1):
            for tile in range(n_tiles1):
                # Compare channels
                for chan in range(min(n_chan1, n_chan2)):
                    for pol in [0, 1, 6, 7]:  # XX,XX and YY,YY
                        val1 = arr1_new[t, tile, chan, pol]
                        val2 = arr2_new[t, tile, chan, pol]
                        if np.isnan(val1) and np.isnan(val2):
                            # This is ok
                            matches += 1
                            # print(f"Match at time {t}, tile {tile}, chan {chan}, pol {pol}: both NaN")
                            continue
                        if not np.isclose(val1, val2, rtol=1e-8, atol=1e-12):
                            print(
                                f"Mismatch at time {t}, tile {tile}, chan {chan}, pol {pol}: {arr1[t, tile, (chan * chan_ratio) : (chan * chan_ratio) + chan_ratio, pol]} => {val1:.4} != {val2:.4}"
                            )
                            mismatches += 1
                        else:
                            # print(f"Match at time {t}, tile {tile}, chan {chan}, pol {pol}: {val1} == {val2}")
                            matches += 1

    print("\nAsymmetrical comparison summary:")
    print(f"  Total comparisons: {matches + mismatches}")
    print(f"  Matches:          {matches}")
    print(f"  Mismatches:       {mismatches}")


def compare_data_stream(fp1, fp2, bytes_to_compare, atol=0.0, rtol=0.0, chunk_elems=131072):
    """
    Compare float64 data streams element-wise in chunks.
    Returns a dict with results:
      - total_elems
      - mismatched_count
      - first_mismatch_index (or None)
      - max_abs_diff
      - max_rel_diff
    Uses tolerance: |a-b| <= atol + rtol * |b| (NumPy-like).
    """
    # Move file pointers to the start of the data region
    fp1.seek(HEADER_TOTAL_LEN)
    fp2.seek(HEADER_TOTAL_LEN)

    elem_size = 8
    total_elems = bytes_to_compare // elem_size

    mismatched_count = 0
    first_mismatch_index = None
    max_abs_diff = 0.0
    max_rel_diff = 0.0

    remaining_elems = total_elems
    idx_base = 0

    # Precompute struct format for chunk unpacking
    # We'll unpack little-endian float64s; format string length depends on chunk size.
    while remaining_elems > 0:
        this_chunk = min(remaining_elems, chunk_elems)
        bytes_needed = this_chunk * elem_size

        b1 = fp1.read(bytes_needed)
        b2 = fp2.read(bytes_needed)

        if len(b1) != bytes_needed or len(b2) != bytes_needed:
            raise ValueError("Unexpected EOF while reading data chunks.")

        # Unpack arrays
        fmt = "<" + ("d" * this_chunk)
        a = struct.unpack(fmt, b1)
        b = struct.unpack(fmt, b2)

        # Compare element-wise
        for i in range(this_chunk):
            ai = a[i]
            bi = b[i]
            abs_diff = abs(ai - bi)
            rel_ref = abs(bi)
            rel_thresh = atol + rtol * rel_ref
            if abs_diff > rel_thresh:
                mismatched_count += 1
                if first_mismatch_index is None:
                    first_mismatch_index = idx_base + i
            if abs_diff > max_abs_diff:
                max_abs_diff = abs_diff
            # Avoid division by zero; track "NumPy-like" rel diff as abs_diff / (|bi| if |bi|>0 else 1)
            rel_diff = abs_diff / (rel_ref if rel_ref > 0.0 else 1.0)
            if rel_diff > max_rel_diff:
                max_rel_diff = rel_diff

        remaining_elems -= this_chunk
        idx_base += this_chunk

    return {
        "total_elems": total_elems,
        "mismatched_count": mismatched_count,
        "first_mismatch_index": first_mismatch_index,
        "max_abs_diff": max_abs_diff,
        "max_rel_diff": max_rel_diff,
    }


def main():
    parser = argparse.ArgumentParser(description="Compare two MWAOCAL-format binary files.")
    parser.add_argument("file1", help="Path to first file")
    parser.add_argument("file2", help="Path to second file")
    parser.add_argument("--atol", type=float, default=0.0, help="Absolute tolerance (default: 0.0)")
    parser.add_argument("--rtol", type=float, default=0.0, help="Relative tolerance (default: 0.0)")
    parser.add_argument(
        "--chunk", type=int, default=131072, help="Number of float64 elements per chunk (default: 131072)"
    )
    args = parser.parse_args()

    # Open files
    with open(args.file1, "rb") as fp1, open(args.file2, "rb") as fp2:
        # Validate file sizes are at least header length
        size1 = os.fstat(fp1.fileno()).st_size
        size2 = os.fstat(fp2.fileno()).st_size
        if size1 < HEADER_TOTAL_LEN or size2 < HEADER_TOTAL_LEN:
            raise ValueError("One or both files are smaller than the minimal header size.")

        # Read headers
        meta1 = read_header(fp1)
        meta2 = read_header(fp2)

        # Report metadata
        print("File 1 metadata:")
        for k, v in meta1.items():
            print(f"  {k}: {v}")
        print("File 2 metadata:")
        for k, v in meta2.items():
            print(f"  {k}: {v}")

        # Compare metadata
        diffs = compare_headers(meta1, meta2)
        if diffs:
            print("\nMetadata differences found:")
            for k, v1, v2 in diffs:
                print(f"  {k}: {v1} != {v2}")
        else:
            print("\nMetadata match.")

        # Compute expected data sizes
        data_bytes_1 = expected_data_bytes(meta1)
        data_bytes_2 = expected_data_bytes(meta2)

        # Validate file sizes vs expected layout
        expected_total_1 = HEADER_TOTAL_LEN + data_bytes_1
        expected_total_2 = HEADER_TOTAL_LEN + data_bytes_2

        if size1 != expected_total_1:
            raise ValueError(f"File 1 size mismatch: expected {expected_total_1} bytes, found {size1} bytes.")
        if size2 != expected_total_2:
            raise ValueError(f"File 2 size mismatch: expected {expected_total_2} bytes, found {size2} bytes.")

        # If metadata differs, we can still attempt data comparison only if the data sizes match
        if data_bytes_1 != data_bytes_2:
            # raise ValueError("Data section sizes differ; cannot compare element-wise.")
            print("Data section sizes differ; cannot compare element-wise.")

            # Attempt asymmetrical comparison
            compare_asymmetrical_data_stream(
                meta1,
                meta2,
                fp1,
                fp2,
                data_bytes_1,
                data_bytes_2,
            )
        else:
            print(f"\nComparing data stream: {data_bytes_1 // 8} float64 elements...")
            result = compare_data_stream(fp1, fp2, data_bytes_1, atol=args.atol, rtol=args.rtol, chunk_elems=args.chunk)

            print("\nData comparison summary:")
            print(f"  Total elements:      {result['total_elems']}")
            print(f"  Mismatched count:    {result['mismatched_count']}")
            print(f"  First mismatch index:{result['first_mismatch_index']}")
            print(f"  Max abs diff:        {result['max_abs_diff']:.6g}")
            print(f"  Max rel diff:        {result['max_rel_diff']:.6g}")

            if result["mismatched_count"] == 0 and not diffs:
                print("\n✅ Files are identical within the given tolerances and metadata matches.")
            elif result["mismatched_count"] == 0 and diffs:
                print("\n⚠️ Data matches within tolerances, but metadata differs (see above).")
            else:
                print("\n❌ Data differs beyond tolerances.")


if __name__ == "__main__":
    main()