"""
Pre-build large fov_pixels PSF grids for forced (outside-FOV) saturated-star fits.

The default PSF grid is fov_pixels=512 (256 px radius).  Diffraction spikes of
out-of-FOV stars extend ~40" along the linear image axes — that's 1290 px in
NIRCam SW (0.031"/pix) and 635 px in NIRCam LW (0.063"/pix).  The standard
grid is too small for those fits.

This script builds large grids — fov_pixels=2048 (SW, radius 1024 px ≈ 32") and
fov_pixels=1024 (LW, radius 512 px ≈ 32") — and saves them next to the existing
small grids using the standard naming convention so `get_psf` can find them.

Usage (one job per (filter, detector); array index selects which):
    python build_large_psf_grids.py <array_index>
"""

import os
import sys

import stpsf

# (filter, detector, fov_pixels)
# detector convention matches stpsf: NRCB1-4 for SW, NRCB5 for LW (NRCBLONG)
TASKS = [
    # SW: F187N + F210M on nrcb2/3/4 for Sickle
    ("F187N", "NRCB2", 2048),
    ("F187N", "NRCB3", 2048),
    ("F187N", "NRCB4", 2048),
    ("F210M", "NRCB2", 2048),
    ("F210M", "NRCB3", 2048),
    ("F210M", "NRCB4", 2048),
    # LW: F335M + F470N + F480M on nrcb5
    ("F335M", "NRCB5", 1024),
    ("F470N", "NRCB5", 1024),
    ("F480M", "NRCB5", 1024),
    # nrcb1 SW (added 2026-06-13: f210m nrcb1 large grid was missing, dropping
    # all forced-seed nrcb1 frames -> nrcb1-quadrant NaN holes in the mosaic)
    ("F187N", "NRCB1", 2048),
    ("F210M", "NRCB1", 2048),
]

OUT_DIR = "/orange/adamginsburg/jwst/sickle/psfs"
NPSF = 16
OVERSAMPLE = 2
# Sickle observation date for OPD selection
OBSDATE = "2024-08-23T12:03:43"


def build(filtername, detector, fov_pixels):
    fn = (f"{OUT_DIR}/nircam_{detector.lower()}_{filtername.lower()}"
          f"_fovp{fov_pixels}_samp{OVERSAMPLE}_npsf{NPSF}.fits")
    if os.path.exists(fn):
        print(f"[skip] {fn} already exists")
        return fn
    print(f"[build] {filtername} {detector} fov_pixels={fov_pixels}")
    psfgen = stpsf.NIRCam()
    psfgen.filter = filtername
    psfgen.detector = detector
    try:
        psfgen.load_wss_opd_by_date(OBSDATE)
    except Exception as exc:
        print(f"  load_wss_opd_by_date warning: {exc}")
    grid = psfgen.psf_grid(num_psfs=NPSF, oversample=OVERSAMPLE,
                           all_detectors=False, fov_pixels=fov_pixels,
                           outdir=OUT_DIR, save=True, outfile=None,
                           overwrite=True)
    print(f"[done] wrote {fn}")
    return fn


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <array_index in [0,{len(TASKS)-1}]>", file=sys.stderr)
        sys.exit(2)
    idx = int(sys.argv[1])
    if idx < 0 or idx >= len(TASKS):
        print(f"index {idx} out of range [0,{len(TASKS)-1}]", file=sys.stderr)
        sys.exit(2)
    filt, det, fov = TASKS[idx]
    build(filt, det, fov)