#!/usr/bin/env python
"""
Test: Computing residuals with and without proper background subtraction.
If we subtract the same background that the fitter used, do the residuals look good?
"""

import numpy as np
from pathlib import Path
from astropy.table import Table
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.stats import sigma_clipped_stats
from photutils.background import LocalBackground
from photutils.psf import PSFPhotometry
from stpsf.utils import to_griddedpsfmodel
import sys
sys.path.insert(0, '/orange/adamginsburg/repos/brick-jwst-2221')

from brick2221.analysis.overfitting_experiment_f480m import (
    load_fits_bundle, cutout_slices,
    replace_nan_pixels_for_fitting
)

# Load data
science_image = Path('/orange/adamginsburg/jwst/sickle/F480M/pipeline/jw03958-o007_t001_nircam_clear-f480m-nrcb_i2d.fits')
stpsf_grid_file = Path('/orange/adamginsburg/jwst/sickle/psfs/nircam_nrcb5_f480m_fovp512_samp2_npsf16.fits')

sci_data, sci_wcs, sci_err, sci_dq, sci_wht = load_fits_bundle(science_image)
psf_model = to_griddedpsfmodel(str(stpsf_grid_file))
fwhm_pix = 2.574

# Load experiment results
exp_outdir = Path('/orange/adamginsburg/jwst/sickle/overfitting_experiments/test_smaller_fit')
stars_tbl = Table.read(exp_outdir / 'cutout_selected_stars.ecsv')

star_id = 0
star_row = stars_tbl[star_id]
xc = float(star_row['xpix'])
yc = float(star_row['ypix'])

halfsize = 18
ysl, xsl = cutout_slices(xc, yc, halfsize, sci_data.shape)
sci_cut = np.asarray(sci_data[ysl, xsl], dtype=float)
sci_fit_cut = replace_nan_pixels_for_fitting(sci_cut, fwhm_pix=fwhm_pix)
sci_err_cut = np.asarray(sci_err[ysl, xsl], dtype=float)

x0 = xc - xsl.start
y0 = yc - ysl.start

from astropy.stats import mad_std
local_noise = mad_std(sci_fit_cut[np.isfinite(sci_fit_cut)], ignore_nan=True)
if not np.isfinite(local_noise) or local_noise <= 0:
    local_noise = 1.0

flux0 = np.nansum(sci_fit_cut[sci_fit_cut > 0]) / 10
init_tbl = Table()
init_tbl['x_0'] = [x0]
init_tbl['y_0'] = [y0]
init_tbl['flux_0'] = [flux0]

# Fit
fit_shape = (7, 7)
localbkg = LocalBackground(6, 10)
uniform_err = np.ones_like(sci_fit_cut)

phot = PSFPhotometry(
    finder=None,
    localbkg_estimator=localbkg,
    psf_model=psf_model,
    fitter=LevMarLSQFitter(),
    fit_shape=fit_shape,
    aperture_radius=2.0 * fwhm_pix,
    progress_bar=False,
)

result = phot(sci_fit_cut, init_params=init_tbl, error=uniform_err)
xfit = float(result['x_fit'][0]) if 'x_fit' in result.colnames else float(result['x_0'][0])
yfit = float(result['y_fit'][0]) if 'y_fit' in result.colnames else float(result['y_0'][0])
flux_fit = float(result['flux_fit'][0]) if 'flux_fit' in result.colnames else np.nan

print(f"Star {star_id}: flux_fit={flux_fit:.2f}")
print()

# Build model manually
psf_eval = psf_model.evaluate(
    x=np.arange(sci_fit_cut.shape[1], dtype=float),
    y=np.arange(sci_fit_cut.shape[0], dtype=float)[:, np.newaxis],
    flux=1.0,
    x_0=xfit,
    y_0=yfit,
)
model = flux_fit * psf_eval

# Estimate background as the fitter would
yy, xx = np.indices(sci_fit_cut.shape, dtype=float)
rr = np.hypot(xx - xfit, yy - yfit)
annulus_mask = (rr >= 6.0) & (rr <= 10.0)

if np.sum(annulus_mask) > 0:
    annulus_data = sci_fit_cut[annulus_mask]
    annulus_data_finite = annulus_data[np.isfinite(annulus_data)]
    bkg_med, _, _ = sigma_clipped_stats(annulus_data_finite, sigma=3.0)
else:
    bkg_med = 0.0

print("="*70)
print("RESIDUALS WITH DIFFERENT BACKGROUND TREATMENTS")
print("="*70)
print(f"Local background estimate: {bkg_med:.4f}\n")

# Method 1: Current approach (WRONG - no background subtraction in residual)
resid_wrong = sci_fit_cut - model
center_idx = (int(np.rint(yfit)), int(np.rint(xfit)))
center_resid_wrong = float(resid_wrong[center_idx])
core_resid_wrong = np.nanmedian(resid_wrong[rr <= 2.0]) if np.sum(rr <= 2.0) > 0 else np.nan

print("Method 1: residual = data - model (CURRENT, NO BKG SUBTRACTION)")
print(f"  Center residual: {center_resid_wrong:.4f}")
print(f"  Core median residual: {core_resid_wrong:.4f}")
print(f"  This matches the overfitting we've been seeing! ❌")

# Method 2: CORRECT - subtract background from data before comparing
data_bkg_sub = sci_fit_cut - bkg_med
resid_correct = data_bkg_sub - model
center_resid_correct = float(resid_correct[center_idx])
core_resid_correct = np.nanmedian(resid_correct[rr <= 2.0]) if np.sum(rr <= 2.0) > 0 else np.nan

print("\nMethod 2: residual = (data - bkg) - model (CORRECT)")
print(f"  Center residual: {center_resid_correct:.4f}")
print(f"  Core median residual: {core_resid_correct:.4f}")
print(f"  Does this look better? ✓")

# Difference
diff = resid_correct - resid_wrong
print(f"\nDifference (correct - wrong): {np.median(diff):.4f} ≈ background value")

# Also test include_localbkg=True
print("\n" + "="*70)
print("COMPARING TO include_localbkg=True")
print("="*70)

model_with_bkg = phot.make_model_image(sci_fit_cut.shape, psf_shape=(21, 21), include_localbkg=True)
resid_with_bkg = sci_fit_cut - model_with_bkg
center_resid_with_bkg = float(resid_with_bkg[center_idx])
core_resid_with_bkg = np.nanmedian(resid_with_bkg[rr <= 2.0]) if np.sum(rr <= 2.0) > 0 else np.nan

print("Method 3: residual = data - (model + localbkg)")
print(f"  Center residual: {center_resid_with_bkg:.4f}")
print(f"  Core median residual: {core_resid_with_bkg:.4f}")

# Summary
print("\n" + "="*70)
print("SUMMARY")
print("="*70)
print(f"{'Method':<45} {'Center Resid':<15} {'Core Median':<15}")
print("-" * 75)
print(f"{'Current (no bkg sub in resid)':<45} {center_resid_wrong:<15.4f} {core_resid_wrong:<15.4f}")
print(f"{'CORRECT (data-bkg)-model':<45} {center_resid_correct:<15.4f} {core_resid_correct:<15.4f}")
print(f"{'include_localbkg=True':<45} {center_resid_with_bkg:<15.4f} {core_resid_with_bkg:<15.4f}")

if abs(center_resid_correct) < abs(center_resid_wrong):
    print(f"\n✓ The CORRECT method eliminates the negative residual bias!")
else:
    print(f"\n⚠ Even the correct method shows residuals")

print("\nDone.")