#!/usr/bin/env python3
"""Compare released bit-packed HSR validity rows with raw finite masks."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--raw-root", required=True, type=Path)
    parser.add_argument("--data-root", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--timestamps", nargs="*", default=None)
    args = parser.parse_args()

    mask_dir = args.data_root / "hsr_valid_mask"
    with (mask_dir / "hsr_valid_mask_meta.json").open("r", encoding="utf-8") as stream:
        meta = json.load(stream)
    timestamps = np.load(mask_dir / "hsr_valid_mask_timestamps.npy", allow_pickle=False).astype("U12")
    mm = np.memmap(
        mask_dir / "hsr_valid_mask.dat",
        dtype=np.uint8,
        mode="r",
        shape=(int(meta["row_count"]), *tuple(meta["row_shape"])),
    )
    lookup = {str(timestamp): index for index, timestamp in enumerate(timestamps)}
    selected = args.timestamps or [str(timestamps[0]), str(timestamps[len(timestamps) // 2]), str(timestamps[-1])]
    checks = []
    for timestamp in selected:
        index = lookup[timestamp]
        raw_path = args.raw_root / timestamp[:8] / f"concat_gk2a_radar_{timestamp}.npy"
        raw = np.load(raw_path, allow_pickle=True).item()
        expected = np.isfinite(np.asarray(raw["hsr"])).reshape(-1)
        actual = np.unpackbits(
            np.asarray(mm[index]).reshape(-1),
            bitorder=str(meta["bitorder"]),
            count=expected.size,
        ).astype(bool)
        differing = int(np.count_nonzero(actual != expected))
        checks.append({"timestamp": timestamp, "pixels": int(expected.size), "differing_pixels": differing})
        if differing:
            raise AssertionError(f"HSR mask differs at {timestamp}: {differing} pixels")

    result = {
        "encoding": "numpy.packbits",
        "bitorder": str(meta["bitorder"]),
        "original_shape": list(meta["original_shape"]),
        "checks": checks,
        "passed": True,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_name(f".{args.output.name}.partial")
    with temporary.open("w", encoding="utf-8") as stream:
        json.dump(result, stream, indent=2)
        stream.write("\n")
    os.replace(temporary, args.output)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()