#!/usr/bin/env python3
"""Bundle the annotated images for the app's validation section.

    python scripts/make_validation_set.py --data-dir /path/to/TIFs_and_CVAT_zips --out validation

Every annotated image is written as a lossless PNG named after the image name
recorded inside its CVAT export (so that cellquant.annotations.discover pairs them
again) next to a copy of that export. Byte-identical duplicate TIFs collapse to one
entry, exactly as in training.
"""
from __future__ import annotations

import argparse
import os
import shutil
import sys

from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from cellquant.annotations import discover, load_image  # noqa: E402
from cellquant.labels import CELL_LABELS, NOSTRIL_LABELS  # noqa: E402


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--out", default=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "validation"))
    a = ap.parse_args(argv)
    samples, _ = discover(a.data_dir, verbose=False)
    os.makedirs(a.out, exist_ok=True)
    for s in samples:
        img = load_image(s.image_path)
        stem = os.path.splitext(os.path.basename(s.annotation.image_name))[0].strip()
        png = os.path.join(a.out, stem + ".png")
        Image.fromarray(img).save(png, optimize=True)
        shutil.copyfile(s.annotation_path, os.path.join(a.out, os.path.basename(s.annotation_path)))
        n_nuc = sum(s.annotation.count(c) for c in CELL_LABELS)
        n_pun = sum(s.annotation.count(c) for c in NOSTRIL_LABELS)
        print(f"{s.key}: {os.path.basename(png)} ({os.path.getsize(png)/1e6:.1f} MB, {n_nuc} nuclei, {n_pun} puncta) + {os.path.basename(s.annotation_path)}")
    check, _ = discover(a.out, verbose=False)
    assert [c.key for c in check] == [s.key for s in samples], "re-discovery of the bundle failed"
    print(f"{len(check)} annotated images in {a.out}")


if __name__ == "__main__":
    main()