#!/usr/bin/env python
"""
Check if background varies radially in oversubtracted stars.
Plot the background as a function of radius to see if sidelobe contamination is present.
"""

import numpy as np
from pathlib import Path
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.stats import sigma_clipped_stats
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import regions
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, cutout_slices,
    read_point_regions
)

# Load data
science_image = Path('/orange/adamginsburg/jwst/sickle/F480M/pipeline/jw03958-o007_t001_nircam_clear-f480m-nrcb_i2d.fits')
region_file = Path('/orange/adamginsburg/jwst/sickle/regions_/diagnostic_oversubtracted_stars_bigger.reg')

print("Loading data...")
sci_data, sci_wcs, sci_err, sci_dq, sci_wht = load_fits_bundle(science_image)

# Read hand-selected regions
point_regions = read_point_regions(region_file)
print(f"Loaded {len(point_regions)} hand-selected oversubtracted stars\n")

# Convert regions to pixel coordinates
region_coords = SkyCoord(
    ra=np.array([reg.center.ra.to_value(u.deg) for reg in point_regions]) * u.deg,
    dec=np.array([reg.center.dec.to_value(u.deg) for reg in point_regions]) * u.deg,
)
xpix_list = []
ypix_list = []
for sc in region_coords:
    x, y = sci_wcs.world_to_pixel(sc)
    xpix_list.append(x)
    ypix_list.append(y)

# Extract radial background profile across all stars
halfsize = 15
max_radius = 16

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

# Store background measurements by radius
all_radial_bkgs = {r: [] for r in r_centers}

for idx, (xc, yc) in enumerate(zip(xpix_list, ypix_list)):
    if idx % 10 == 0:
        print(f"Processing star {idx}/{len(point_regions)}...", flush=True)

    # Skip if near edge
    if xc < halfsize or xc > sci_data.shape[1] - halfsize or \
       yc < halfsize or yc > sci_data.shape[0] - halfsize:
        continue

    # Extract cutout
    ysl, xsl = cutout_slices(xc, yc, halfsize, sci_data.shape)
    sci_cut = np.asarray(sci_data[ysl, xsl], dtype=float)

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

    # Measure background in radial annuli
    yy, xx = np.indices(sci_cut.shape, dtype=float)
    rr = np.hypot(xx - x0, yy - y0)

    for i in range(len(r_centers)):
        r_in = r_edges[i]
        r_out = r_edges[i + 1]
        annulus_mask = (rr >= r_in) & (rr < r_out)

        if np.sum(annulus_mask) < 3:
            continue

        annulus_data = sci_cut[annulus_mask]
        annulus_data_finite = annulus_data[np.isfinite(annulus_data)]

        if len(annulus_data_finite) < 3:
            continue

        # Sigma-clipped median
        med, _, _ = sigma_clipped_stats(annulus_data_finite, sigma=3.0)
        all_radial_bkgs[r_centers[i]].append(float(med))

# Compute statistics by radius
print("\nRadial Background Profile")
print(f"{'Radius':<10} {'Median':<12} {'Mean':<12} {'Std':<12} {'N_samples':<10}")
print("-" * 50)

radii_with_data = []
medians = []
stds = []

for r in r_centers:
    if len(all_radial_bkgs[r]) == 0:
        continue
    vals = np.array(all_radial_bkgs[r])
    med = np.median(vals)
    mn = np.mean(vals)
    std = np.std(vals)
    n = len(vals)
    print(f"{r:<10.1f} {med:<12.4f} {mn:<12.4f} {std:<12.4f} {n:<10d}")
    radii_with_data.append(r)
    medians.append(med)
    stds.append(std)

# Plot
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))

# Plot 1: Background vs radius
ax1.errorbar(radii_with_data, medians, yerr=stds, fmt='o-', capsize=5, markersize=6)
ax1.axhline(y=0, color='k', linestyle='--', alpha=0.3, label='Zero')
# Mark LocalBackground regions
ax1.axvspan(2, 5, alpha=0.1, color='red', label='LocalBkg(2,5)')
ax1.axvspan(6, 10, alpha=0.1, color='green', label='LocalBkg(6,10) [current]')
ax1.axvspan(5, 15, alpha=0.1, color='blue', label='LocalBkg(5,15)')
ax1.set_xlabel('Radius [pix]')
ax1.set_ylabel('Background [counts]')
ax1.set_title('Radial Background Profile')
ax1.grid(True, alpha=0.3)
ax1.legend(fontsize=8)

# Plot 2: Just the median values to see the trend
ax2.plot(radii_with_data, medians, 'o-', markersize=6, linewidth=2)
ax2.axhline(y=0, color='k', linestyle='--', alpha=0.3)
ax2.axvline(x=2, color='red', linestyle='--', alpha=0.3, label='PSF FWHM ~2.6 pix')
ax2.set_xlabel('Radius [pix]')
ax2.set_ylabel('Background Median [counts]')
ax2.set_title('Background Trend with Radius')
ax2.grid(True, alpha=0.3)
ax2.legend()

fig.suptitle('Radial Background Profile in Hand-Selected Oversubtracted Stars')
fig.tight_layout()
fig.savefig('/orange/adamginsburg/jwst/sickle/radial_background_profile.png', dpi=120)
print(f"\nPlot saved to /orange/adamginsburg/jwst/sickle/radial_background_profile.png")

# Analysis
print("\n" + "="*70)
print("INTERPRETATION")
print("="*70)
if len(radii_with_data) > 1:
    inner_r = np.mean(medians[:3]) if len(medians) > 2 else medians[0]
    outer_r = np.mean(medians[-3:]) if len(medians) > 2 else medians[-1]
    diff = inner_r - outer_r
    pct_change = 100 * diff / outer_r if outer_r != 0 else 0

    print(f"\nInner background (r~1-2):   {inner_r:.4f}")
    print(f"Outer background (r~14-16): {outer_r:.4f}")
    print(f"Difference:                 {diff:.4f}")
    print(f"Percent change:             {pct_change:.1f}%")

    if pct_change > 50:
        print("\n⚠ SIGNIFICANT RADIAL VARIATION: Inner regions much brighter than outer.")
        print("  This suggests PSF sidelobe contamination in inner annuli.")
    elif pct_change > 10:
        print("\n⚠ MODERATE RADIAL VARIATION: Some background structure present.")
    else:
        print("\n✓ Background is relatively constant with radius.")

print("\nDone.")