"""Offline numeric check of the pinned GOOSE decoder on retained Worker RGB.

Runs in the existing pinned image, without changing its runner or checkpoint.
This diagnoses the adapter; it is not semantic accuracy or navigation acceptance.
"""

import argparse
import csv
import hashlib
import importlib.util
import json
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    spec = importlib.util.spec_from_file_location("reference", "/assets/ddrnet-goose-runner.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    model, _, _ = module.load_model("ddrnet", Path("/assets/ddrnet-checkpoint.pth"))
    tensor, _ = module.preprocess(Image.open(args.image))
    tensor = tensor.cuda()
    with torch.inference_mode():
        logits = module.logits_from_output(model(tensor)).float()
        legacy = torch.sigmoid(logits).argmax(1)
        direct = logits.argmax(1)
        saturated = (torch.sigmoid(logits) == 1).sum(1)
    names = {}
    with open("/assets/ddrnet-goose-mapping.csv") as stream:
        names = {int(r["label_key"]): r["class_name"] for r in csv.DictReader(stream)}
    args.output.mkdir(exist_ok=True, parents=True)
    report = {
        "source_sha256": hashlib.sha256(args.image.read_bytes()).hexdigest(),
        "monotonic_ns": time.monotonic_ns(),
        "logit_range": [float(logits.min()), float(logits.max())],
        "changed_pixels": int((legacy != direct).sum()),
        "saturated_tie_pixels": int((saturated > 1).sum()),
    }
    for name, mask in (("reference", legacy), ("direct", direct)):
        mask = mask[0].cpu().numpy().astype(np.uint8)
        Image.fromarray(mask).save(args.output / (name + ".png"))
        ids, counts = np.unique(mask, return_counts=True)
        report[name] = {names[int(i)]: int(counts[index]) for index, i in enumerate(ids)}
    (args.output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report))


if __name__ == "__main__":
    main()