#!/usr/bin/env python3
"""
5×4 grid showing photometry cataloging progress in the benchmark_cutout region.

Rows : Original i2d | iter0 residual | iter2 residual | iter3 residual
Cols : F187N | F210M | F335M | F470N | F480M

Normalization: unified asinh(1%–99%) per column so rows are directly comparable.
Colormap     : gray_r (reversed grayscale).

Outputs go to  catalog/diagnostic_images/  and follow the pattern:
  benchmark_progress_{images|catalog}_{basic|iterative}_{noinfill|infill}.png

That gives 8 files total (2 catalog variants × 2 methods × 2 infill variants).
"""

import warnings
import glob
import os

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from astropy.io import fits
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.table import Table, vstack
from astropy.nddata import Cutout2D
from astropy.visualization import (AsinhStretch, AsymmetricPercentileInterval,
                                   ImageNormalize)

# ── paths ────────────────────────────────────────────────────────────────────
BASE    = '/orange/adamginsburg/jwst/sickle'
OUTDIR  = os.path.join(BASE, 'catalog', 'diagnostic_images')
os.makedirs(OUTDIR, exist_ok=True)

# ── region ───────────────────────────────────────────────────────────────────
REGION_RA  = 266.568138784   # deg
REGION_DEC = -28.804072264   # deg
REGION_W   = 3.3566          # arcsec
REGION_H   = 3.3953          # arcsec
CENTER     = SkyCoord(ra=REGION_RA * u.deg, dec=REGION_DEC * u.deg)
CUTOUT_SIZE = (REGION_H * u.arcsec, REGION_W * u.arcsec)   # (ny, nx)

# ── filters ──────────────────────────────────────────────────────────────────
FILTERS     = ['f187n', 'f210m', 'f335m', 'f470n', 'f480m']
FILT_LABELS = ['F187N', 'F210M', 'F335M', 'F470N', 'F480M']
FILT_DIRS   = {f: os.path.join(BASE, f.upper()) for f in FILTERS}

# ── grid row definitions ─────────────────────────────────────────────────────
# (iteration_token, stage_label, row_title)
ROWS = [
    (None,    'i2d',   'Original i2d'),
    (None,    'iter0', 'iter0 residual'),
    ('iter2', 'iter2', 'iter2 residual'),
    ('iter3', 'iter3', 'iter3 residual'),
]

# ── catalog overlay colours ───────────────────────────────────────────────────
CAT_COLORS = {'iter0': '#00AAFF', 'iter2': '#FF4400', 'iter3': '#00EE00'}


# ═══════════════════════════════════════════════════════════════════════════════
# Path helpers
# ═══════════════════════════════════════════════════════════════════════════════

def i2d_orig_path(filt):
    # Use the pipeline-generated full mosaic i2d so all filters share the same
    # WCS as the residual mosaics and have a consistent (not over-subtracted)
    # background.  The mastDownload f470n product was heavily bg-subtracted.
    return os.path.join(FILT_DIRS[filt], 'pipeline',
                        f'jw03958-o007_t001_nircam_clear-{filt}-nrcb_i2d.fits')


def residual_i2d_path(filt, iteration, method, infilled):
    """
    iteration : None  → iter0 (no label), 'iter2', 'iter3'
    method    : 'basic' | 'iterative'
    infilled  : bool
    """
    fdir   = FILT_DIRS[filt]
    prefix = f'jw03958-o007_t001_nircam_clear-{filt}-nrcb'
    iter_tok = f'_{iteration}' if iteration else ''
    infill   = '_infilled'    if infilled   else ''
    fname  = f'{prefix}{iter_tok}_daophot_{method}_residual{infill}_i2d.fits'
    return os.path.join(fdir, 'pipeline', fname)


# ═══════════════════════════════════════════════════════════════════════════════
# FITS loading
# ═══════════════════════════════════════════════════════════════════════════════

def load_cutout(path):
    """Return (data_2d, wcs) cutout for the benchmark region, or (None, None)."""
    if not os.path.exists(path):
        return None, None
    try:
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            hdul = fits.open(path, ignore_missing_end=True, memmap=False)
        ext_names = [e.name for e in hdul]
        ext = 'SCI' if 'SCI' in ext_names else 1
        hdr  = hdul[ext].header
        data = hdul[ext].data.astype(float)
        hdul.close()
        wcs = WCS(hdr, naxis=2)
        cut = Cutout2D(data, CENTER, CUTOUT_SIZE, wcs=wcs,
                       mode='partial', fill_value=np.nan)
        return cut.data, cut.wcs
    except Exception as exc:
        print(f'  [warn] {os.path.basename(path)}: {exc}')
        return None, None


# ═══════════════════════════════════════════════════════════════════════════════
# Per-column normalization
# ═══════════════════════════════════════════════════════════════════════════════

def column_norms(cutout_grid, vmin_fixed=None, vmax_fixed=None):
    """
    cutout_grid : list of lists  [row][col] → ndarray or None
    Returns list of ImageNormalize (one per column).

    If vmin_fixed / vmax_fixed are given those values are used directly for
    every column (fully fixed scale, identical across all rows AND columns).
    Otherwise vmin/vmax are derived from the first (i2d) row of each column
    using the 1%–99% asinh interval; the same norm is then reused for every
    row in that column so all panels share an identical colour scale.
    """
    ncols = len(cutout_grid[0])
    norms = []
    for j in range(ncols):
        if vmin_fixed is not None and vmax_fixed is not None:
            norm = ImageNormalize(vmin=vmin_fixed, vmax=vmax_fixed,
                                  stretch=AsinhStretch())
        else:
            d0 = cutout_grid[0][j]
            finite = d0[np.isfinite(d0)] if d0 is not None else np.array([])
            if finite.size:
                vmin, vmax = AsymmetricPercentileInterval(1, 99).get_limits(finite)
                norm = ImageNormalize(vmin=vmin, vmax=vmax, stretch=AsinhStretch())
            else:
                norm = ImageNormalize(vmin=0, vmax=1, stretch=AsinhStretch())
        norms.append(norm)
    return norms


# ═══════════════════════════════════════════════════════════════════════════════
# Catalog loaders
# ═══════════════════════════════════════════════════════════════════════════════

def _box_mask(sc, half_w_deg, half_h_deg):
    dra  = (sc.ra.deg  - CENTER.ra.deg) * np.cos(np.deg2rad(CENTER.dec.deg))
    ddec =  sc.dec.deg - CENTER.dec.deg
    return (np.abs(dra) < half_w_deg) & (np.abs(ddec) < half_h_deg)


_HW = ((REGION_W / 2) / 3600, (REGION_H / 2) / 3600)   # (half_w_deg, half_h_deg)


def _sc_from_table(t):
    for col in ('skycoord_avg', 'skycoord', 'skycoord_centroid'):
        if col in t.colnames:
            return SkyCoord(t[col])
    if 'ra' in t.colnames:
        return SkyCoord(ra=t['ra'] * u.deg, dec=t['dec'] * u.deg)
    return None


def load_iter0_catalog(filt):
    # Prefer _allcols which carries skycoord_avg (astrometry-corrected mean).
    # Fall back to the slim file if allcols doesn't exist yet.
    for suffix in ('_allcols.fits', '.fits'):
        path = os.path.join(BASE, 'catalogs',
                            f'{filt}_merged_indivexp_merged_dao_basic{suffix}')
        if os.path.exists(path):
            break
    else:
        return None
    sc = _sc_from_table(Table.read(path))
    return None if sc is None else sc[_box_mask(sc, *_HW)]


def load_iter2_catalog(filt):
    pattern = os.path.join(FILT_DIRS[filt],
                           f'{filt}_nrcb*_iter2_daophot_basic.fits')
    tables = []
    for f in sorted(glob.glob(pattern)):
        try:
            tables.append(Table.read(f))
        except Exception:
            pass
    if not tables:
        return None
    sc = _sc_from_table(vstack(tables, metadata_conflicts='silent'))
    return None if sc is None else sc[_box_mask(sc, *_HW)]


def load_iter3_catalog(filt):
    path = os.path.join(BASE, 'benchmark', 'photutils_iter3', 'catalogs',
                        f'{filt}_photutils_iter3_nrcb.fits')
    if not os.path.exists(path):
        return None
    sc = _sc_from_table(Table.read(path))
    return None if sc is None else sc[_box_mask(sc, *_HW)]


# ═══════════════════════════════════════════════════════════════════════════════
# Main grid builder
# ═══════════════════════════════════════════════════════════════════════════════

def _overlay_catalog(ax, sc_in_box, wcs, data_shape, color, marker='o', s=14, lw=1.0):
    if sc_in_box is None or len(sc_in_box) == 0:
        return
    try:
        px, py = wcs.world_to_pixel(sc_in_box)
        ny, nx = data_shape
        ok = (px >= 0) & (px < nx) & (py >= 0) & (py < ny)
        if ok.sum():
            ax.scatter(px[ok], py[ok], s=s, marker=marker,
                       linewidths=lw, facecolors='none',
                       edgecolors=color, alpha=0.85, zorder=5)
    except Exception as exc:
        print(f'    [warn] overlay: {exc}')


def build_grid(method, infilled, with_catalog, out_path,
               vmin_fixed=None, vmax_fixed=None):
    """
    method      : 'basic' | 'iterative'
    infilled    : bool
    with_catalog: bool
    out_path    : output PNG path
    vmin_fixed / vmax_fixed : if given, all panels use this fixed range
    """
    nrows, ncols = len(ROWS), len(FILTERS)

    # ── 1. load all cutouts ──────────────────────────────────────────────────
    cutouts = [[None] * ncols for _ in range(nrows)]
    wcses   = [[None] * ncols for _ in range(nrows)]

    for j, filt in enumerate(FILTERS):
        for i, (iteration, stage, _) in enumerate(ROWS):
            if stage == 'i2d':
                path = i2d_orig_path(filt)
            else:
                path = residual_i2d_path(filt, iteration, method, infilled)
            d, w = load_cutout(path)
            cutouts[i][j] = d
            wcses[i][j]   = w

    # ── 2. per-column norms ──────────────────────────────────────────────────
    col_norm = column_norms(cutouts, vmin_fixed=vmin_fixed, vmax_fixed=vmax_fixed)

    # ── 3. pre-load catalogs once if needed ──────────────────────────────────
    cats = {}
    if with_catalog:
        for filt in FILTERS:
            cats[filt] = {
                'iter0': load_iter0_catalog(filt),
                'iter2': load_iter2_catalog(filt),
                'iter3': load_iter3_catalog(filt),
            }

    # ── 4. build figure ──────────────────────────────────────────────────────
    fig_w = ncols * 2.6 + 0.7
    fig_h = nrows * 2.6 + 0.9
    fig, axes = plt.subplots(nrows, ncols, figsize=(fig_w, fig_h), squeeze=False)
    fig.subplots_adjust(left=0.07, right=0.99, top=0.93, bottom=0.05,
                        wspace=0.03, hspace=0.05)

    for j, filt_label in enumerate(FILT_LABELS):
        axes[0, j].set_title(filt_label, fontsize=11, fontweight='bold', pad=3)

    for i, (_, _, row_title) in enumerate(ROWS):
        axes[i, 0].set_ylabel(row_title, fontsize=9, labelpad=3)

    for i, (iteration, stage, _) in enumerate(ROWS):
        for j, filt in enumerate(FILTERS):
            ax   = axes[i, j]
            data = cutouts[i][j]
            wcs  = wcses[i][j]
            norm = col_norm[j]

            if data is None or not np.any(np.isfinite(data)):
                ax.text(0.5, 0.5, 'N/A', ha='center', va='center',
                        transform=ax.transAxes, color='gray', fontsize=9)
            else:
                ax.imshow(data, origin='lower', norm=norm, cmap='gray_r',
                          interpolation='nearest')

                if with_catalog and wcs is not None:
                    # Each successive foreground layer is drawn smaller and thinner
                    # so overlapping detections produce a visible bullseye pattern.
                    _SIZES = [22, 13, 6]   # background → foreground
                    _LWS   = [1.2, 0.8, 0.5]
                    if stage == 'i2d':
                        # all three catalogs on the original image;
                        # draw iter3 first (bottom/largest), iter0 last (top/thinnest)
                        for k, st in enumerate(['iter3', 'iter2', 'iter0']):
                            _overlay_catalog(ax, cats[filt][st], wcs,
                                             data.shape, CAT_COLORS[st],
                                             s=_SIZES[k], lw=_LWS[k])
                    else:
                        # cumulative: show all stages up to and including this one;
                        # earlier catalogs are background (larger), later are foreground (thinner)
                        for k, st in enumerate(['iter0', 'iter2', 'iter3']):
                            _overlay_catalog(ax, cats[filt][st], wcs,
                                             data.shape, CAT_COLORS[st],
                                             s=_SIZES[k], lw=_LWS[k])
                            if st == stage:
                                break

            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_aspect('equal')

    if with_catalog:
        handles = [mpatches.Patch(color=CAT_COLORS[s], label=f'{s} detections')
                   for s in ('iter0', 'iter2', 'iter3')]
        fig.legend(handles=handles, loc='lower center', ncol=3,
                   fontsize=9, framealpha=0.9, bbox_to_anchor=(0.5, 0.0))

    infill_str = 'infill' if infilled else 'noinfill'
    if vmin_fixed is not None:
        scale_str = f'asinh  vmin={vmin_fixed}  vmax={vmax_fixed}  (fixed, all panels)'
    else:
        scale_str = 'asinh 1–99% from i2d row, identical per column'
    fig.suptitle(
        f'Sickle benchmark-cutout  |  method={method}  infill={infill_str}\n'
        rf'RA=266.568°, Dec=−28.804°, ~7″×7″  |  gray_r, {scale_str}',
        fontsize=10)

    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    print(f'  → {out_path}')


# ═══════════════════════════════════════════════════════════════════════════════
# Entry point — generate all 8 variants
# ═══════════════════════════════════════════════════════════════════════════════

if __name__ == '__main__':
    # Fixed display range: vmin=-1, vmax=100 (image units, typically MJy/sr).
    # Identical for every panel so rows and columns are all directly comparable.
    VMIN, VMAX = -1, 100

    for method in ('basic', 'iterative'):
        for infilled in (False, True):
            infill_tag = 'infill' if infilled else 'noinfill'
            for with_catalog in (False, True):
                cat_tag = 'catalog' if with_catalog else 'images'
                fname   = f'benchmark_progress_{cat_tag}_{method}_{infill_tag}.png'
                out     = os.path.join(OUTDIR, fname)
                print(f'Building {fname} …')
                build_grid(method=method, infilled=infilled,
                           with_catalog=with_catalog, out_path=out,
                           vmin_fixed=VMIN, vmax_fixed=VMAX)
    print('Done.')