#!/usr/bin/env python
"""
Inspect fitter's fit_info to check:
1. Does the fitter claim good convergence (low chi^2, small residuals)?
2. Do the returned parameters match the fit_info?
3. If we manually compute chi^2, do we get the same value?
"""

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, 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[star_id]
xc = float(star_row['xpix'])
yc = float(star_row['ypix'])

# 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

# Initial params
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"Star {star_id} at ({xc:.1f}, {yc:.1f})")
print(f"Cutout center: ({x0:.2f}, {y0:.2f})")
print(f"Initial flux: {flux0:.2f}\n")

# Fit with LocalBackground(6,10) and fit_shape=(7,7) - the problematic one
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,
)

print("Running photometry...")
result = phot(sci_fit_cut, init_params=init_tbl, error=uniform_err)

# Extract fit parameters
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"Fitted params: x={xfit:.4f}, y={yfit:.4f}, flux={flux_fit:.2f}\n")

# Access the fitter's fit_info
print("="*70)
print("FITTER'S FIT_INFO")
print("="*70)
if hasattr(phot, 'fitter') and hasattr(phot.fitter, 'fit_info'):
    fit_info = phot.fitter.fit_info
    if fit_info:
        print(f"ierr: {fit_info.get('ierr', 'N/A')}")
        print(f"nfev: {fit_info.get('nfev', 'N/A')} (function evaluations)")
        print(f"message: {fit_info.get('message', 'N/A')}")

        # fvec contains residuals at each data point
        fvec = fit_info.get('fvec', None)
        if fvec is not None:
            fvec = np.asarray(fvec)
            chi2_from_fvec = float(np.sum(fvec**2))
            print(f"\nResiduals (fvec):")
            print(f"  Shape: {fvec.shape}")
            print(f"  Min: {np.min(fvec):.6f}")
            print(f"  Max: {np.max(fvec):.6f}")
            print(f"  Mean: {np.mean(fvec):.6f}")
            print(f"  Std: {np.std(fvec):.6f}")
            print(f"  Chi^2 (sum of fvec^2): {chi2_from_fvec:.6g}")

        # cov_x is the covariance matrix
        cov_x = fit_info.get('cov_x', None)
        if cov_x is not None:
            cov_x = np.asarray(cov_x)
            print(f"\nCovariance matrix:")
            print(f"  Shape: {cov_x.shape}")
            print(f"  Diagonal (variances): {np.diag(cov_x)}")
            print(f"  Std errors: {np.sqrt(np.diag(cov_x))}")

print("\n" + "="*70)
print("MANUALLY COMPUTING CHI^2")
print("="*70)

# Manually evaluate PSF at fitted position
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

# Method 1: Without background subtraction
resid_no_bkg = sci_fit_cut - model
chi2_no_bkg = np.nansum(resid_no_bkg**2)  # Assuming uniform weight = 1
print(f"Without background subtraction:")
print(f"  Chi^2: {chi2_no_bkg:.6g}")
print(f"  RMS residual: {np.sqrt(np.mean(resid_no_bkg[np.isfinite(resid_no_bkg)]**2)):.6f}")

# Method 2: With background subtraction in fit window
# Estimate background from annulus as the fitter would
yy, xx = np.indices(sci_fit_cut.shape, dtype=float)
rr = np.hypot(xx - xfit, yy - yfit)

# Define fit window and annulus
fit_hw = fit_shape[0] // 2
fit_mask = (np.abs(xx - xfit) <= fit_hw) & (np.abs(yy - yfit) <= fit_hw)

annulus_inner = 6.0
annulus_outer = 10.0
annulus_mask = (rr >= annulus_inner) & (rr <= annulus_outer)

if np.sum(annulus_mask) > 0:
    from astropy.stats import sigma_clipped_stats
    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)
    print(f"\nEstimated local background: {bkg_med:.4f}")

    # Data with background subtracted (in fit window only)
    sci_bkg_sub = np.array(sci_fit_cut, copy=True)
    sci_bkg_sub[fit_mask] -= bkg_med

    model_bkg_sub = np.array(model, copy=True)
    # Don't subtract background from model - it's just PSF

    resid_with_bkg = sci_bkg_sub - model_bkg_sub
    # Only compute chi^2 in fit window
    resid_fit_window = resid_with_bkg[fit_mask]
    chi2_with_bkg = np.nansum(resid_fit_window**2)

    print(f"\nWith background subtraction in fit window:")
    print(f"  Fit window pixels: {np.sum(fit_mask)}")
    print(f"  Chi^2 (fit window): {chi2_with_bkg:.6g}")
    print(f"  RMS residual (fit window): {np.sqrt(np.mean(resid_fit_window[np.isfinite(resid_fit_window)]**2)):.6f}")

print("\n" + "="*70)
print("FINAL MODEL AND RESIDUALS")
print("="*70)

print(f"Data value at center ({int(np.rint(yfit))}, {int(np.rint(xfit))}): {sci_fit_cut[int(np.rint(yfit)), int(np.rint(xfit))]:.4f}")
print(f"Model value at center: {model[int(np.rint(yfit)), int(np.rint(xfit))]:.4f}")
print(f"Residual at center: {resid_no_bkg[int(np.rint(yfit)), int(np.rint(xfit))]:.4f}")

print("\nDone.")