#!/usr/bin/env python
"""
Test hypothesis: LocalBackground is UNDERESTIMATING the true background.
Compare LocalBackground estimate to the true background from far outer annulus.
"""

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 regions
import sys

sys.path.insert(0, '/orange/adamginsburg/repos/brick-jwst-2221')

from brick2221.analysis.overfitting_experiment_f480m import (
    load_fits_bundle, 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,
)

halfsize = 15

# Definitions of background regions
configs = [
    ("True BG (r=14-16)", 14, 16),  # Far outer: "true" background
    ("LocalBkg(6,10)", 6, 10),       # Current production
    ("LocalBkg(5,15)", 5, 15),
    ("LocalBkg(2,5)", 2, 5),         # Sidelobe contaminated
]

print(f"Testing {len(point_regions)} hand-selected oversubtracted stars")
print()

results = {name: [] for name, _, _ in configs}

for idx, sc in enumerate(region_coords):
    if idx % 10 == 0:
        print(f"Processing star {idx}/{len(point_regions)}...", flush=True)

    xc, yc = sci_wcs.world_to_pixel(sc)

    # 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 each annulus
    yy, xx = np.indices(sci_cut.shape, dtype=float)
    rr = np.hypot(xx - x0, yy - y0)

    for name, r_in, r_out in configs:
        annulus_mask = (rr >= r_in) & (rr <= r_out)

        if np.sum(annulus_mask) < 3:
            results[name].append(np.nan)
            continue

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

        if len(annulus_data_finite) < 3:
            results[name].append(np.nan)
            continue

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

# Analyze
print("\n" + "="*70)
print("BACKGROUND UNDERESTIMATION TEST")
print("="*70)

for name in results:
    results[name] = np.array(results[name])
    valid = ~np.isnan(results[name])
    if np.sum(valid) > 0:
        vals = results[name][valid]
        print(f"\n{name:20}")
        print(f"  N measured: {np.sum(valid)}")
        print(f"  Median:     {np.median(vals):>10.4f}")
        print(f"  Mean:       {np.mean(vals):>10.4f}")
        print(f"  Std:        {np.std(vals):>10.4f}")

# Compare to true background
print("\n" + "="*70)
print("UNDERESTIMATION vs TRUE BACKGROUND (r=14-16)")
print("="*70)

true_bkg = results["True BG (r=14-16)"]
valid_true = ~np.isnan(true_bkg)

if np.sum(valid_true) > 0:
    for name, _, _ in configs:
        if name == "True BG (r=14-16)":
            continue

        est_bkg = results[name]
        both_valid = valid_true & (~np.isnan(est_bkg))

        if np.sum(both_valid) == 0:
            continue

        true_vals = true_bkg[both_valid]
        est_vals = est_bkg[both_valid]
        diff = est_vals - true_vals
        pct_error = 100 * diff / np.abs(true_vals)

        print(f"\n{name}:")
        print(f"  Estimated median:   {np.median(est_vals):>10.4f}")
        print(f"  True background:    {np.median(true_vals):>10.4f}")
        print(f"  Difference (est-true): {np.median(diff):>10.4f}")
        print(f"  Percent error:      {np.median(pct_error):>10.1f}%")

        if np.median(diff) > 5:
            print(f"  ✓ OVERESTIMATING background (conservative)")
        elif np.median(diff) < -5:
            print(f"  ✗ UNDERESTIMATING background (dangerous!)")
        else:
            print(f"  ≈ Roughly unbiased")

print("\nDone.")