#!/usr/bin/env python
"""
Debug script to check if error weighting is biasing the fit.
Tests if central pixels are being down-weighted relative to outer pixels.
"""

import numpy as np
from pathlib import Path
from astropy.table import Table
from astropy.visualization import simple_norm
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

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)
crowd_wht_map = compute_crowdsource_weight_map(sci_data, sci_err, dq=sci_dq, wht=sci_wht)

# 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')

# Pick 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_err_cut = np.asarray(sci_err[ysl, xsl], dtype=float)
weight_cut = np.asarray(crowd_wht_map[ysl, xsl], dtype=float)

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

print(f"\nCutout shape: {sci_cut.shape}")
print(f"Star center in cutout: ({x0:.2f}, {y0:.2f})")

# Analyze error/weight distribution
print(f"\n=== Error Map Statistics ===")
print(f"sci_err_cut:")
print(f"  min={np.nanmin(sci_err_cut):.6e}, max={np.nanmax(sci_err_cut):.6e}")
print(f"  median={np.nanmedian(sci_err_cut):.6e}, mean={np.nanmean(sci_err_cut):.6e}")

print(f"\nweight_cut (1/err²):")
print(f"  min={np.nanmin(weight_cut):.6e}, max={np.nanmax(weight_cut):.6e}")
print(f"  median={np.nanmedian(weight_cut):.6e}, mean={np.nanmean(weight_cut):.6e}")

# Compute weights by radius from center
yy, xx = np.indices(sci_cut.shape, dtype=float)
rr = np.hypot(xx - x0, yy - y0)

# Radial binning
r_edges = np.arange(0, 20, 1)
r_centers = 0.5 * (r_edges[:-1] + r_edges[1:])

print(f"\n=== Weighting by Radius ===")
print(f"{'Radius':<10} {'Data(med)':<15} {'Error(med)':<15} {'Weight(med)':<15} {'N_pixels':<10}")
print("-" * 65)

for i in range(len(r_centers)):
    sel = (rr >= r_edges[i]) & (rr < r_edges[i+1])
    n = np.sum(sel)
    if n == 0:
        continue

    data_med = np.nanmedian(sci_cut[sel])
    err_med = np.nanmedian(sci_err_cut[sel])
    wt_med = np.nanmedian(weight_cut[sel])

    print(f"{r_centers[i]:<10.1f} {data_med:<15.2f} {err_med:<15.2e} {wt_med:<15.2e} {n:<10d}")

# Analyze central vs outer regions
core_r = 2.0
outer_r_in = 3.0
outer_r_out = 6.0

core_sel = rr <= core_r
outer_sel = (rr >= outer_r_in) & (rr <= outer_r_out)

print(f"\n=== Core vs Outer Weighting ===")
print(f"Core (r<={core_r}):")
print(f"  N_pixels={np.sum(core_sel)}")
print(f"  data_median={np.nanmedian(sci_cut[core_sel]):.2f}")
print(f"  error_median={np.nanmedian(sci_err_cut[core_sel]):.2e}")
print(f"  weight_median={np.nanmedian(weight_cut[core_sel]):.2e}")
print(f"  weight_mean={np.nanmean(weight_cut[core_sel]):.2e}")

print(f"\nOuter (r∈[{outer_r_in},{outer_r_out}]):")
print(f"  N_pixels={np.sum(outer_sel)}")
print(f"  data_median={np.nanmedian(sci_cut[outer_sel]):.2f}")
print(f"  error_median={np.nanmedian(sci_err_cut[outer_sel]):.2e}")
print(f"  weight_median={np.nanmedian(weight_cut[outer_sel]):.2e}")
print(f"  weight_mean={np.nanmean(weight_cut[outer_sel]):.2e}")

# Check if outer pixels are weighted higher relative to their signal
print(f"\n=== Signal-to-Weight Analysis ===")
core_snr = np.nanmedian(sci_cut[core_sel]) / np.nanmedian(sci_err_cut[core_sel])
outer_snr = np.nanmedian(sci_cut[outer_sel]) / np.nanmedian(sci_err_cut[outer_sel])

print(f"Core SNR: {core_snr:.2f}")
print(f"Outer SNR: {outer_snr:.2f}")

core_weighted_snr = core_snr * np.sqrt(np.nanmedian(weight_cut[core_sel]))
outer_weighted_snr = outer_snr * np.sqrt(np.nanmedian(weight_cut[outer_sel]))

print(f"Core weighted SNR: {core_weighted_snr:.2f}")
print(f"Outer weighted SNR: {outer_weighted_snr:.2f}")

# Create diagnostic plot
fig, axes = plt.subplots(2, 3, figsize=(14, 8))

# Data
norm_data = simple_norm(sci_cut, stretch='log', percent=99.5)
im0 = axes[0, 0].imshow(sci_cut, origin='lower', cmap='gray', norm=norm_data)
axes[0, 0].set_title('Science Data')
axes[0, 0].plot([x0], [y0], 'r+', ms=10, mew=2)
plt.colorbar(im0, ax=axes[0, 0])

# Error map
im1 = axes[0, 1].imshow(sci_err_cut, origin='lower', cmap='viridis')
axes[0, 1].set_title('Error Map')
axes[0, 1].plot([x0], [y0], 'r+', ms=10, mew=2)
plt.colorbar(im1, ax=axes[0, 1])

# Weight map
im2 = axes[0, 2].imshow(weight_cut, origin='lower', cmap='hot')
axes[0, 2].set_title('Weight Map (1/err²)')
axes[0, 2].plot([x0], [y0], 'r+', ms=10, mew=2)
plt.colorbar(im2, ax=axes[0, 2])

# Radius mask
im3 = axes[1, 0].imshow(rr, origin='lower', cmap='twilight')
axes[1, 0].set_title('Radius from Center')
circ_core = plt.Circle((x0, y0), core_r, fill=False, edgecolor='red', lw=1.5, label=f'core r<{core_r}')
circ_outer_in = plt.Circle((x0, y0), outer_r_in, fill=False, edgecolor='cyan', lw=1.5, ls='--', label=f'outer ring')
circ_outer_out = plt.Circle((x0, y0), outer_r_out, fill=False, edgecolor='cyan', lw=1.5)
axes[1, 0].add_patch(circ_core)
axes[1, 0].add_patch(circ_outer_in)
axes[1, 0].add_patch(circ_outer_out)
axes[1, 0].legend(fontsize=8)
plt.colorbar(im3, ax=axes[1, 0])

# Radial profile: data vs error
ax = axes[1, 1]
for i in range(len(r_centers)):
    sel = (rr >= r_edges[i]) & (rr < r_edges[i+1])
    if np.sum(sel) == 0:
        continue
    data_med = np.nanmedian(sci_cut[sel])
    err_med = np.nanmedian(sci_err_cut[sel])
    ax.plot(r_centers[i], data_med, 'bo', ms=6, label='data' if i == 0 else '')
    ax.plot(r_centers[i], err_med, 'rx', ms=6, label='error' if i == 0 else '')
ax.set_xlabel('Radius [pix]')
ax.set_ylabel('Median Value')
ax.set_title('Data vs Error by Radius')
ax.legend()
ax.grid(True, alpha=0.3)

# Radial profile: weight
ax = axes[1, 2]
for i in range(len(r_centers)):
    sel = (rr >= r_edges[i]) & (rr < r_edges[i+1])
    if np.sum(sel) == 0:
        continue
    wt_med = np.nanmedian(weight_cut[sel])
    ax.plot(r_centers[i], wt_med, 'go', ms=6)
ax.set_xlabel('Radius [pix]')
ax.set_ylabel('Median Weight')
ax.set_title('Weight by Radius')
ax.grid(True, alpha=0.3)

fig.suptitle(f'Star {star_id} Weighting Analysis')
fig.tight_layout()
fig.savefig('/orange/adamginsburg/jwst/sickle/debug_weighting.png', dpi=120)
print(f"\nDiagnostic plot saved to /orange/adamginsburg/jwst/sickle/debug_weighting.png")

print("\nDone.")