"""Step 13: blind contact sheets for visual vetting of the f>=0.75 candidates and 40 random controls.

Each object gets a wide image (4.3 arcmin, with a 1.75 arcmin-radius circle = half the Arecibo beam FWHM)
and a zoom (77 arcsec).  Objects are shown in a seeded random order under anonymous codes so that the
classifier does not know which are candidates and which are controls; the key is data/visual_key.csv.
The classifications themselves were entered by hand into data/visual_classification.csv (columns code, uid,
vis_class, note); that file is an INPUT of s14 and is not regenerated by the pipeline.
Outputs: RAW/vetting_sheets/sheet_XX.png (regenerable), data/visual_key.csv
"""
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from common import RAW, DATA

PER_SHEET = 9
SEED = 4242


def best_jpg(tag, uid):
    for suffix in ("", "_ls-dr10", "_sdss", "_ps1"):
        p = os.path.join(RAW, "cutouts", f"{tag}_{uid}{suffix}.jpg")
        if os.path.exists(p) and os.path.getsize(p) > 6000:
            return p, (suffix or "_ls-dr9")[1:]
    return None, "none"


def main():
    t = pd.read_csv(os.path.join(DATA, "vet_targets.csv"))
    v = t[t.visual].copy()
    rng = np.random.default_rng(SEED)
    v = v.iloc[rng.permutation(len(v))].reset_index(drop=True)
    v["code"] = [f"V{k + 1:03d}" for k in range(len(v))]
    outdir = os.path.join(RAW, "vetting_sheets")   # ~2 MB each: kept out of the shared folder
    os.makedirs(outdir, exist_ok=True)
    layers = []
    for s0 in range(0, len(v), PER_SHEET):
        chunk = v.iloc[s0:s0 + PER_SHEET]
        fig, axes = plt.subplots(3, 6, figsize=(15.4, 7.9), dpi=100)
        for ax in axes.ravel():
            ax.axis("off")
        for k, r in enumerate(chunk.itertuples()):
            row, col = divmod(k, 3)
            lay = "none"
            for j, tag in enumerate(("wide", "zoom")):
                ax = axes[row, 2 * col + j]
                p, lay_j = best_jpg(tag, r.uid)
                if p is None:
                    ax.text(0.5, 0.5, "no image", ha="center", va="center", transform=ax.transAxes)
                    continue
                im = plt.imread(p)
                ax.imshow(im)
                if tag == "wide":
                    lay = lay_j
                    ax.add_patch(Circle((im.shape[1] / 2, im.shape[0] / 2), 105, fill=False, color="yellow", lw=0.8, ls="--"))
                    ax.set_title(f"{r.code}  wide 4.3'", fontsize=9, pad=2)
                else:
                    ax.set_title(f"{r.code}  zoom 77\"", fontsize=9, pad=2)
            layers.append(lay)
        fig.subplots_adjust(left=0.005, right=0.995, top=0.965, bottom=0.005, wspace=0.02, hspace=0.10)
        fig.savefig(os.path.join(outdir, f"sheet_{s0 // PER_SHEET + 1:02d}.png"))
        plt.close(fig)
    v["image_layer"] = layers
    v[["code", "uid", "role", "image_layer"]].to_csv(os.path.join(DATA, "visual_key.csv"), index=False)
    print(f"  contact sheets: {len(v)} objects on {int(np.ceil(len(v) / PER_SHEET))} sheets; "
          f"image layers: {v.image_layer.value_counts().to_dict()}")


if __name__ == "__main__":
    main()