"""Optional helper (not a pipeline stage): Legacy Survey image cutouts of the unresolved three-way candidates
flagged at f >= 0.5 (final cuts), for a visual check for blends, galaxies and artefacts.

    python3 cutouts.py            # writes $DIM_STARS_RAW/cutouts/<source_id>_<layer>.jpg and a contact sheet

The visual inspection itself is manual; its outcome is reported in the text as such and resolves nothing.
"""
import os
import urllib.request

import pandas as pd

import dimlib as L
from common import DATA_MAIN, RAW, log

URL = "https://www.legacysurvey.org/viewer/cutout.jpg?ra=%.6f&dec=%.6f&layer=%s&pixscale=0.262&size=256"
LAYERS = ("ls-dr9", "ls-dr10", "unwise-neo7")      # DR10 covers the southern stars outside DR9


def run():
    out = os.path.join(RAW, "cutouts")
    os.makedirs(out, exist_ok=True)
    u = pd.read_csv(os.path.join(DATA_MAIN, "candidates_unresolved.csv"))
    keep = pd.Series(False, index=u.index)
    for f in (0.5, 0.75, 0.9):
        keep |= u["flag_%g" % f] & (u["status_%g" % f] == "unresolved")
    u = u[keep].drop_duplicates("source_id").sort_values("D")
    files = []
    for _, r in u.iterrows():
        for layer in LAYERS:
            p = os.path.join(out, "%d_%s.jpg" % (r.source_id, layer))
            if not os.path.exists(p):
                try:
                    with urllib.request.urlopen(URL % (r.ra, r.dec, layer), timeout=60) as h, open(p, "wb") as fh:
                        fh.write(h.read())
                except Exception as e:
                    log("cutout failed for %d %s: %r" % (r.source_id, layer, e))
                    continue
            files.append((r.source_id, layer, p))
    log("cutouts: %d files for %d stars in %s" % (len(files), len(u), out))
    # contact sheet (optical layer), 67 arcsec on a side, north up
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        # optical layer: DR9 where it has data, otherwise DR10 (an empty cutout is a small uniform JPEG)
        opt = []
        for i in u.source_id:
            cand = [p for j, lay, p in files if j == i and lay.startswith("ls-") and os.path.getsize(p) > 3000]
            if cand:
                opt.append((i, cand[0]))
        n = len(opt)
        ncol = 6
        nrow = (n + ncol - 1) // ncol
        fig, axes = plt.subplots(nrow, ncol, figsize=(2.2 * ncol, 2.4 * nrow))
        for ax in axes.ravel():
            ax.axis("off")
        for ax, (i, p) in zip(axes.ravel(), opt):
            try:
                ax.imshow(plt.imread(p))
            except Exception:
                continue
            ax.set_title(str(i), fontsize=7)
        fig.tight_layout()
        fig.savefig(os.path.join(out, "contact_sheet_optical.png"), dpi=110)
        log("optical images for %d of %d stars" % (len(opt), len(u)))
    except Exception as e:
        log("contact sheet failed: %r" % e)


if __name__ == "__main__":
    run()