#!/usr/bin/env python
"""
Test: Does the initial flux estimate affect the fitted result?
Try different initial flux values to see if we get different minima.
"""

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
)

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

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)

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

# Compute "true" flux from entire cutout minus background
yy, xx = np.indices(sci_fit_cut.shape, dtype=float)
rr = np.hypot(xx - x0, yy - y0)

# Estimate background from far outer annulus
annulus_mask = (rr >= 14) & (rr <= 16)
if np.sum(annulus_mask) > 0:
    annulus_data = sci_fit_cut[annulus_mask]
    annulus_data_finite = annulus_data[np.isfinite(annulus_data)]
    true_bkg, _, _ = sigma_clipped_stats(annulus_data_finite, sigma=3.0)
else:
    true_bkg = 0.0

# Estimate "true" flux as sum of all pixels above background
flux_true = np.nansum(np.clip(sci_fit_cut - true_bkg, 0, None))

localbkg = LocalBackground(6, 10)
uniform_err = np.ones_like(sci_fit_cut)

print(f"Star {star_id}")
print(f"Background (outer annulus r=14-16): {true_bkg:.2f}")
print(f"'True' flux (sum of core above bkg): {flux_true:.2f}\n")

# Test different initial flux estimates
flux_estimates = [
    (flux_true * 0.25, "25% of true"),
    (flux_true * 0.5, "50% of true"),
    (flux_true * 1.0, "100% of true (full core)"),
    (flux_true * 2.0, "200% of true"),
    (flux_true * 4.0, "400% of true"),
]

results = []

for flux0, label in flux_estimates:
    init_tbl = Table()
    init_tbl['x_0'] = [x0]
    init_tbl['y_0'] = [y0]
    init_tbl['flux_0'] = [flux0]

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

    result = phot(sci_fit_cut, init_params=init_tbl, error=uniform_err)
    flux_fit = float(result['flux_fit'][0])

    # Compute residual
    xfit = float(result['x_fit'][0])
    yfit = float(result['y_fit'][0])

    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 local background for this fit
    rr_fit = np.hypot(xx - xfit, yy - yfit)
    annulus_mask_fit = (rr_fit >= 6.0) & (rr_fit <= 10.0)
    if np.sum(annulus_mask_fit) > 0:
        annulus_data = sci_fit_cut[annulus_mask_fit]
        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

    data_bkg_sub = sci_fit_cut - bkg_med
    resid = data_bkg_sub - model

    center_idx = (int(np.rint(yfit)), int(np.rint(xfit)))
    center_resid = float(resid[center_idx])

    # Get chi^2 from fitter
    if hasattr(phot, 'fitter') and hasattr(phot.fitter, 'fit_info') and phot.fitter.fit_info:
        fvec = np.asarray(phot.fitter.fit_info.get('fvec', []))
        chi2_reported = float(np.sum(fvec**2))
    else:
        chi2_reported = np.nan

    results.append({
        'label': label,
        'flux0': flux0,
        'flux_fit': flux_fit,
        'center_resid': center_resid,
        'chi2': chi2_reported,
    })

print(f"{'Initial Flux':<25} {'Initial':<12} {'Fitted':<12} {'Center Resid':<14} {'Chi^2':<12}")
print("-" * 80)

for res in results:
    print(f"{res['label']:<25} {res['flux0']:<12.0f} {res['flux_fit']:<12.0f} {res['center_resid']:<14.4f} {res['chi2']:<12.0f}")

print("\nAnalysis:")
fitted_fluxes = [r['flux_fit'] for r in results]
chi2s = [r['chi2'] for r in results]

print(f"All fitted fluxes: min={min(fitted_fluxes):.0f}, max={max(fitted_fluxes):.0f}")
print(f"Range: {max(fitted_fluxes) - min(fitted_fluxes):.0f}")

if min(chi2s) == max(chi2s):
    print("✓ All initial fluxes converge to the SAME chi^2!")
else:
    print(f"⚠ Different chi^2 values: min={min(chi2s):.0f}, max={max(chi2s):.0f}")

print("\nDone.")