#!/usr/bin/env python
"""
Test if uniform weights vs error-weighted fits produce the same overfitting.
"""

import numpy as np
from pathlib import Path
from astropy.table import Table
from astropy.modeling.fitting import LevMarLSQFitter
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, load_fits_data_and_wcs,
    compute_crowdsource_weight_map, 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')

print("Loading data...")
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 overfitting experiment results
exp_outdir = Path('/orange/adamginsburg/jwst/sickle/overfitting_experiments/test_smaller_fit')
stars_tbl = Table.read(exp_outdir / 'cutout_selected_stars.ecsv')

# Same star as before
star_id = 0
star_row = stars_tbl[stars_tbl['star_id'] == star_id][0]
xc = float(star_row['xpix'])
yc = float(star_row['ypix'])

print(f"Analyzing star_id={star_id} at x={xc:.2f}, y={yc:.2f}")

# Extract cutout
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

# Estimate initial flux
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]

print(f"Initial flux estimate: {flux0:.2f}")
print(f"Cutout shape: {sci_fit_cut.shape}\n")

# Test fit_shape (7,7) - the one that shows overfitting
fit_shape = (7, 7)
localbkg = LocalBackground(5, 15)

# Fit 1: Using error weights (as currently done)
print("=== FIT 1: Error-weighted (current approach) ===")
phot1 = 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,
)
result1 = phot1(sci_fit_cut, init_params=init_tbl, error=np.where(np.isfinite(sci_err_cut), sci_err_cut, 1e11))
xfit1 = float(result1['x_fit'][0]) if 'x_fit' in result1.colnames else float(result1['x_0'][0])
yfit1 = float(result1['y_fit'][0]) if 'y_fit' in result1.colnames else float(result1['y_0'][0])
flux1 = float(result1['flux_fit'][0]) if 'flux_fit' in result1.colnames else np.nan

model1 = phot1.make_model_image(sci_fit_cut.shape, psf_shape=(21, 21), include_localbkg=False)
resid1 = sci_fit_cut - model1
center_idx = (int(np.rint(yfit1)), int(np.rint(xfit1)))
center_resid1 = float(resid1[center_idx])

print(f"Fitted flux: {flux1:.2f} (vs initial {flux0:.2f}, delta={flux1-flux0:.2f})")
print(f"Fitted center: ({xfit1:.4f}, {yfit1:.4f})")
print(f"Center residual: {center_resid1:.4f}")

# Fit 2: Using uniform weights
print("\n=== FIT 2: Uniform weights ===")
phot2 = 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,
)
uniform_err = np.ones_like(sci_fit_cut)
result2 = phot2(sci_fit_cut, init_params=init_tbl, error=uniform_err)
xfit2 = float(result2['x_fit'][0]) if 'x_fit' in result2.colnames else float(result2['x_0'][0])
yfit2 = float(result2['y_fit'][0]) if 'y_fit' in result2.colnames else float(result2['y_0'][0])
flux2 = float(result2['flux_fit'][0]) if 'flux_fit' in result2.colnames else np.nan

model2 = phot2.make_model_image(sci_fit_cut.shape, psf_shape=(21, 21), include_localbkg=False)
resid2 = sci_fit_cut - model2
center_resid2 = float(resid2[center_idx])

print(f"Fitted flux: {flux2:.2f} (vs initial {flux0:.2f}, delta={flux2-flux0:.2f})")
print(f"Fitted center: ({xfit2:.4f}, {yfit2:.4f})")
print(f"Center residual: {center_resid2:.4f}")

# Fit 3: Using no weighting (error=None)
print("\n=== FIT 3: No weighting (error=None) ===")
phot3 = 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,
)
result3 = phot3(sci_fit_cut, init_params=init_tbl, error=None)
xfit3 = float(result3['x_fit'][0]) if 'x_fit' in result3.colnames else float(result3['x_0'][0])
yfit3 = float(result3['y_fit'][0]) if 'y_fit' in result3.colnames else float(result3['y_0'][0])
flux3 = float(result3['flux_fit'][0]) if 'flux_fit' in result3.colnames else np.nan

model3 = phot3.make_model_image(sci_fit_cut.shape, psf_shape=(21, 21), include_localbkg=False)
resid3 = sci_fit_cut - model3
center_resid3 = float(resid3[center_idx])

print(f"Fitted flux: {flux3:.2f} (vs initial {flux0:.2f}, delta={flux3-flux0:.2f})")
print(f"Fitted center: ({xfit3:.4f}, {yfit3:.4f})")
print(f"Center residual: {center_resid3:.4f}")

# Compare
print("\n=== COMPARISON ===")
print(f"{'Weighting':<20} {'Flux':<12} {'Flux Delta':<15} {'Center Resid':<15}")
print("-" * 65)
print(f"{'Error-weighted':<20} {flux1:<12.2f} {flux1-flux0:<15.2f} {center_resid1:<15.4f}")
print(f"{'Uniform':<20} {flux2:<12.2f} {flux2-flux0:<15.2f} {center_resid2:<15.4f}")
print(f"{'None (unweighted)':<20} {flux3:<12.2f} {flux3-flux0:<15.2f} {center_resid3:<15.4f}")

if abs(center_resid1 - center_resid2) < 5:
    print(f"\n✓ Uniform weights produce SAME overfitting as error-weighted (diff={abs(center_resid1-center_resid2):.4f})")
else:
    print(f"\n✗ Uniform weights produce DIFFERENT result (diff={abs(center_resid1-center_resid2):.4f})")

print("\nDone.")