#!/usr/bin/env python
"""Small smoke test: crop a 512x512 region of F480M and run crowdsource."""
import sys
from pathlib import Path

import numpy as np
from astropy.io import fits

sys.path.insert(0, str(Path(__file__).resolve().parent))
from run_crowdsource_mosaic import WrappedPSFModel, load_gridded_psf

from crowdsource.crowdsource_base import fit_im

mosaic = "/orange/adamginsburg/jwst/sickle/F480M/pipeline/jw03958-o007_t001_nircam_clear-f480m-nrcb_i2d.fits"
hdul = fits.open(mosaic)
sci = hdul["SCI"].data.astype(np.float32)
print("mosaic shape:", sci.shape)
# 512x512 crop centered-ish where there's data
ny, nx = sci.shape
yc, xc = ny // 2, nx // 2
sub = sci[yc - 256:yc + 256, xc - 256:xc + 256].copy()
err = hdul["ERR"].data.astype(np.float32)[yc - 256:yc + 256, xc - 256:xc + 256]
w = np.zeros_like(err)
good = err > 0
w[good] = 1.0 / err[good]
sub = np.nan_to_num(sub, nan=0.0, posinf=0.0, neginf=0.0)
print("sub median", np.nanmedian(sub), "min", np.nanmin(sub), "max", np.nanmax(sub))

grid = load_gridded_psf("F480M", detector="nrcb5")
psf = WrappedPSFModel(grid, stampsz=19)

print("Running fit_im...")
cat, mod, sky, psffit = fit_im(sub, psf, weights=w, nskyx=0, nskyy=0, refit_psf=False, verbose=True, threshold=5.0)
print(f"Detected {len(cat)} sources in 512x512 crop.")