"""
Single-frame test: how does relaxing the iter2 DAO shape cuts affect
recovery of the by-eye F2550W stars?

This is a STANDALONE driver — it does not run PSF photometry.  It just
runs DAOStarFinder on the (NaN-replaced, satstar-subtracted) data of one
frame using the iter2 threshold (5th-percentile local-noise) and varies
the sharp/round cuts.  Then it cross-matches the daofind output against
the by-eye CRTF and reports recovered counts + how many of the daofind
sources are spurious (no real source within 0.5") for each shape config.

We test 4 shape configurations:
  - tight  : sharplo=0.50, sharphi=1.00, roundlo=-0.3, roundhi=0.3  (current iter2)
  - mid    : sharplo=0.40, sharphi=1.20, roundlo=-0.4, roundhi=0.4
  - loose  : sharplo=0.30, sharphi=1.40, roundlo=-0.5, roundhi=0.5  (matches iter1)
  - veryloose : sharplo=0.10, sharphi=1.80, roundlo=-1.0, roundhi=1.0

Frame chosen: jw05365998001_06101_00003_mirimage_cal.fits — an obs-998
frame with high source density and a few bright Sgr B2 stars in extended
emission, so it's a useful stress test.
"""

import os, glob, sys, re
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy import wcs
from astropy.table import Table
from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.convolution import Gaussian2DKernel, convolve_fft, interpolate_replace_nans
from photutils.detection import DAOStarFinder
import warnings
from astropy.wcs import FITSFixedWarning
warnings.simplefilter('ignore', category=FITSFixedWarning)

sys.path.insert(0, '/orange/adamginsburg/repos/brick-jwst-2221')
from jwst.datamodels import dqflags
from brick2221.analysis.crowdsource_catalogs_long import (
    compute_local_noise_map, _bad_dq_bitmask,
)

PIPE = '/orange/adamginsburg/jwst/sgrb2/F2550W/pipeline'
CRTF = '/home/nbudaiev/orange_link/sgrb2/jwst/NB/sgrb2_jwst/by_eye_catalogs/catalog_by_eye_point_source_f2550w.crtf'
FWHM_PIX = 7.30
MATCH_RADIUS = 0.4 * u.arcsec   # match the validation radius

CAL = f'{PIPE}/jw05365998001_06101_00003_mirimage_cal.fits'
SAT_MODEL = CAL.replace('.fits', '_satstar_model.fits')

# ----- Load the by-eye catalog -----
re_pos = re.compile(r'symbol\s*\[\[\s*([\d.+-]+)\s*deg\s*,\s*([\d.+-]+)\s*deg', re.IGNORECASE)
ra, dec = [], []
with open(CRTF) as f:
    for line in f:
        m = re_pos.search(line)
        if m:
            ra.append(float(m.group(1))); dec.append(float(m.group(2)))
byeye = SkyCoord(ra=np.array(ra)*u.deg, dec=np.array(dec)*u.deg, frame='icrs')
print(f'by-eye catalog: {len(byeye)}')

# ----- Build the same nan_replaced, satstar-subtracted data the pipeline uses -----
with fits.open(CAL) as f:
    sci = f['SCI'].data.astype(np.float32)
    err = f['ERR'].data.astype(np.float32)
    dq = f['DQ'].data
    ww = wcs.WCS(f['SCI'].header)
    ny, nx = sci.shape
print(f'frame shape: {(ny, nx)}')

bad_bm = _bad_dq_bitmask('MIRI')
is_bad = (dq & bad_bm) != 0
nan_data = sci.copy()
nan_data[is_bad] = np.nan
kernel = Gaussian2DKernel(x_stddev=FWHM_PIX/2.355)
nan_replaced = interpolate_replace_nans(nan_data, kernel, convolve=convolve_fft)
# Satstar model subtraction (matches do_photometry_step iter2 path)
if os.path.exists(SAT_MODEL):
    sat_model = fits.getdata(SAT_MODEL).astype(np.float32)
    sat_finite = np.where(np.isfinite(sat_model), sat_model, 0.0)
    nan_replaced = nan_replaced - sat_finite
    print(f'subtracted satstar model ({SAT_MODEL})')

# ----- Threshold using the fixed (5th-percentile, data!=0) recipe -----
lnm = compute_local_noise_map(nan_replaced, smooth_sigma_pix=3.0)
finite = np.isfinite(lnm) & (lnm > 0) & (nan_replaced != 0)
threshold = float(np.nanpercentile(lnm[finite], 5))
print(f'iter2 threshold (5th-pct local noise, data!=0) = {threshold:.3f}')

# ----- Cross-match recovery for each shape config -----
configs = [
    dict(label='tight',    sharplo=0.50, sharphi=1.00, roundlo=-0.3, roundhi=0.3),
    dict(label='mid',      sharplo=0.40, sharphi=1.20, roundlo=-0.4, roundhi=0.4),
    dict(label='loose',    sharplo=0.30, sharphi=1.40, roundlo=-0.5, roundhi=0.5),
    dict(label='veryloose',sharplo=0.10, sharphi=1.80, roundlo=-1.0, roundhi=1.0),
]

# Which by-eye stars actually fall on this frame?
xy = ww.world_to_pixel(byeye)
xpix = np.asarray(xy[0]); ypix = np.asarray(xy[1])
in_frame = (xpix >= 0) & (xpix < nx) & (ypix >= 0) & (ypix < ny)
# also require not on bad DQ
ix_ = np.clip(np.round(xpix).astype(int), 0, nx-1)
iy_ = np.clip(np.round(ypix).astype(int), 0, ny-1)
on_bad = np.zeros(len(byeye), dtype=bool)
on_bad[in_frame] = (dq[iy_[in_frame], ix_[in_frame]] & bad_bm) != 0
covered_mask = in_frame & ~on_bad
sub_byeye = byeye[covered_mask]
print(f'by-eye stars in this frame & not on bad DQ: {covered_mask.sum()} of {len(byeye)}')

results = []
for cfg in configs:
    finder = DAOStarFinder(threshold=threshold, fwhm=FWHM_PIX,
                           sharplo=cfg['sharplo'], sharphi=cfg['sharphi'],
                           roundlo=cfg['roundlo'], roundhi=cfg['roundhi'])
    finds = finder(nan_replaced, mask=is_bad)
    if finds is None:
        n_finds = 0; recovered = 0; spurious = 0
    else:
        n_finds = len(finds)
        find_skyc = ww.pixel_to_world(np.asarray(finds['xcentroid']),
                                      np.asarray(finds['ycentroid']))
        # Recovery: how many sub_byeye stars have a find within MATCH_RADIUS?
        if len(find_skyc) > 0 and len(sub_byeye) > 0:
            idx_, sep_, _ = sub_byeye.match_to_catalog_sky(find_skyc)
            recovered = int((sep_ < MATCH_RADIUS).sum())
            # Spurious: finds with NO sub_byeye within MATCH_RADIUS.  This
            # is a *very* loose definition since by-eye is incomplete; use
            # cautiously.
            idx2_, sep2_, _ = find_skyc.match_to_catalog_sky(sub_byeye)
            spurious = int((sep2_ > MATCH_RADIUS).sum())
        else:
            recovered = spurious = 0
    results.append((cfg['label'], cfg['sharplo'], cfg['sharphi'],
                    cfg['roundlo'], cfg['roundhi'],
                    n_finds, recovered, spurious))
    print(f'  {cfg["label"]:9s}  sharp={cfg["sharplo"]:.2f}-{cfg["sharphi"]:.2f}  '
          f'round=±{cfg["roundhi"]:.2f}  '
          f'finds={n_finds:5d}  recovered={recovered}/{covered_mask.sum()}  '
          f'spurious≈{spurious}')

t = Table(rows=results,
          names=['label', 'sharplo', 'sharphi', 'roundlo', 'roundhi',
                 'n_finds', 'recovered_byeye', 'spurious_finds'])
out = f'{PIPE}/F2550W_iter2_shape_sweep_one_frame.fits'
t.write(out, overwrite=True)
print(f'wrote {out}')
print(f'\nbase: {covered_mask.sum()} byeye stars on this frame & DQ-clean')