#!/usr/bin/env python
"""Color-color & color-magnitude diagrams for wd1 / wd2 / w51 (EWOCS comparison).

Uses the cross-band merged Vega-mag catalogs:
  wd1/wd2 : w51-ewocs/catalogs/{field}_ewocs_merged_vega.fits  (built by
            build_merged_catalog.py -- union cross-match of per-filter vetted DAO)
  w51     : the pipeline multiband merged catalog (already has Vega mags)

All three share the column convention  mag_vega_<f> / emag_ab_<f> / qfit_<f> /
is_saturated_<f>, so one accessor drives every plot.

Quality gate (per filter used, from the sgrc/brick convention):
  finite mag, emag_ab < EMAG_MAX, and qfit < QFIT_MAX for UNSATURATED stars only
  (saturated cores have spuriously inflated qfit -- do not cut them on qfit).
"""
import os
import warnings
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.table import Table

BASE = "/orange/adamginsburg/jwst"
OUT = f"{BASE}/w51-ewocs/figures"

# Prefer the full pipeline cross-band merge if it has landed (jobs 37790299/300);
# fall back to the hand-built union catalog otherwise. Same mag_vega_<f> columns.
_PIPE = ("basic_merged_indivexp_photometry_tables_merged_resbgsub_m7.fits")
CATS = {
    "wd1": (f"{BASE}/wd1/catalogs/{_PIPE}"
            if os.path.exists(f"{BASE}/wd1/catalogs/{_PIPE}")
            else f"{BASE}/w51-ewocs/catalogs/wd1_ewocs_merged_vega.fits"),
    "wd2": (f"{BASE}/wd2/catalogs/{_PIPE}"
            if os.path.exists(f"{BASE}/wd2/catalogs/{_PIPE}")
            else f"{BASE}/w51-ewocs/catalogs/wd2_ewocs_merged_vega.fits"),
    "w51": f"{BASE}/w51/catalogs/"
           "basic_merged_indivexp_photometry_tables_merged_qualcuts_oksep2221.fits",
}
TITLE = {"wd1": "Westerlund 1", "wd2": "Westerlund 2", "w51": "W51"}

FILTERS = {
    "wd1": ["f115w", "f150w", "f164n", "f187n", "f200w", "f212n",
            "f277w", "f323n", "f405n", "f444w", "f466n"],
    "wd2": ["f115w", "f150w", "f162m", "f164n", "f182m", "f187n", "f200w",
            "f212n", "f250m", "f277w", "f300m", "f323n", "f335m", "f405n",
            "f410m", "f444w", "f466n"],
    "w51": ["f140m", "f162m", "f182m", "f187n", "f210m", "f335m", "f360m",
            "f405n", "f410m", "f480m"],
}

QFIT_MAX = 0.4
EMAG_MAX = 0.15


def mag(t, fl):
    return np.asarray(t[f"mag_vega_{fl}"], dtype=float)


def issat(t, fl):
    c = f"is_saturated_{fl}"
    if c not in t.colnames:
        return np.zeros(len(t), bool)
    return np.asarray(np.nan_to_num(np.asarray(t[c], float)), bool) \
        if t[c].dtype != bool else np.asarray(t[c], bool)


def good(t, *filts):
    m = np.ones(len(t), dtype=bool)
    for fl in filts:
        mg = mag(t, fl)
        em = np.asarray(t[f"emag_ab_{fl}"], dtype=float)
        qf = np.asarray(t[f"qfit_{fl}"], dtype=float)
        sat = issat(t, fl)
        m &= np.isfinite(mg) & np.isfinite(em) & (em < EMAG_MAX)
        m &= sat | (np.isfinite(qf) & (qf < QFIT_MAX))
    return m


def load(field):
    t = Table.read(CATS[field])
    return t


# ---------------------------------------------------------------------------
def _hexbin_cmd(ax, t, b, r, y, xlim=None, ylim=None, cmap="viridis"):
    m = good(t, b, r, y)
    color = mag(t, b)[m] - mag(t, r)[m]
    ymag = mag(t, y)[m]
    ok = np.isfinite(color) & np.isfinite(ymag)
    color, ymag = color[ok], ymag[ok]
    n = len(color)
    if n == 0:
        ax.text(0.5, 0.5, "no data", ha="center", va="center",
                transform=ax.transAxes)
        return 0
    if xlim is None:
        xlim = (np.nanpercentile(color, 1), np.nanpercentile(color, 99))
    if ylim is None:
        ylim = (np.nanpercentile(ymag, 0.5), np.nanpercentile(ymag, 99.5))
    ext = (xlim[0], xlim[1], ylim[0], ylim[1])
    if n > 400:
        ax.hexbin(color, ymag, gridsize=110, bins="log", cmap=cmap,
                  mincnt=1, extent=ext)
    else:
        ax.scatter(color, ymag, s=4, alpha=0.5, c="k", edgecolors="none")
    ax.set_xlim(xlim)
    ax.set_ylim(ylim[1], ylim[0])  # inverted (bright at top)
    ax.set_xlabel(f"{b.upper()} - {r.upper()}  (Vega)")
    ax.set_ylabel(f"{y.upper()}  (Vega)")
    return n


def _hexbin_ccd(ax, t, bx, rx, by, ry, xlim=None, ylim=None, cmap="viridis"):
    m = good(t, bx, rx, by, ry)
    cx = mag(t, bx)[m] - mag(t, rx)[m]
    cy = mag(t, by)[m] - mag(t, ry)[m]
    ok = np.isfinite(cx) & np.isfinite(cy)
    cx, cy = cx[ok], cy[ok]
    n = len(cx)
    if n == 0:
        ax.text(0.5, 0.5, "no data", ha="center", va="center",
                transform=ax.transAxes)
        return 0
    if xlim is None:
        xlim = (np.nanpercentile(cx, 1), np.nanpercentile(cx, 99))
    if ylim is None:
        ylim = (np.nanpercentile(cy, 1), np.nanpercentile(cy, 99))
    ext = (xlim[0], xlim[1], ylim[0], ylim[1])
    if n > 400:
        ax.hexbin(cx, cy, gridsize=110, bins="log", cmap=cmap, mincnt=1,
                  extent=ext)
    else:
        ax.scatter(cx, cy, s=4, alpha=0.5, c="k", edgecolors="none")
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.set_xlabel(f"{bx.upper()} - {rx.upper()}  (Vega)")
    ax.set_ylabel(f"{by.upper()} - {ry.upper()}  (Vega)")
    return n


# ---------------------------------------------------------------------------
def per_field_cmd_grid(field, cmds, ncol=4):
    t = load(field)
    n = len(cmds)
    nrow = int(np.ceil(n / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(4.2 * ncol, 4.6 * nrow),
                             squeeze=False)
    for ax, (b, r, y) in zip(axes.flat, cmds):
        nn = _hexbin_cmd(ax, t, b, r, y)
        ax.set_title(f"{b.upper()}-{r.upper()} vs {y.upper()}  (N={nn})",
                     fontsize=9)
    for ax in axes.flat[n:]:
        ax.axis("off")
    fig.suptitle(f"{TITLE[field]} ({field}) - color-magnitude diagrams "
                 f"(Vega; emag<{EMAG_MAX}, qfit<{QFIT_MAX} unsat)", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fn = f"{OUT}/{field}_cmd_grid.png"
    fig.savefig(fn, dpi=125)
    plt.close(fig)
    print("wrote", fn)


def per_field_ccd_grid(field, ccds, ncol=4):
    t = load(field)
    n = len(ccds)
    nrow = int(np.ceil(n / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(4.4 * ncol, 4.4 * nrow),
                             squeeze=False)
    for ax, (bx, rx, by, ry) in zip(axes.flat, ccds):
        nn = _hexbin_ccd(ax, t, bx, rx, by, ry)
        ax.set_title(f"{bx.upper()}-{rx.upper()} vs {by.upper()}-{ry.upper()}"
                     f"  (N={nn})", fontsize=9)
    for ax in axes.flat[n:]:
        ax.axis("off")
    fig.suptitle(f"{TITLE[field]} ({field}) - color-color diagrams "
                 f"(Vega; emag<{EMAG_MAX}, qfit<{QFIT_MAX} unsat)", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fn = f"{OUT}/{field}_ccd_grid.png"
    fig.savefig(fn, dpi=125)
    plt.close(fig)
    print("wrote", fn)


def side_by_side_cmd(fields, b, r, y, tag, xlim=None, ylim=None):
    """One CMD (b-r vs y) across N fields, shared axes, side by side."""
    fig, axes = plt.subplots(1, len(fields), figsize=(5.2 * len(fields), 6),
                             squeeze=False)
    # common ranges from union of fields if not given
    if xlim is None or ylim is None:
        allc, ally = [], []
        for f in fields:
            t = load(f)
            m = good(t, b, r, y)
            allc.append(mag(t, b)[m] - mag(t, r)[m])
            ally.append(mag(t, y)[m])
        allc = np.concatenate(allc)
        ally = np.concatenate(ally)
        allc = allc[np.isfinite(allc)]
        ally = ally[np.isfinite(ally)]
        if xlim is None:
            xlim = (np.nanpercentile(allc, 1), np.nanpercentile(allc, 99))
        if ylim is None:
            ylim = (np.nanpercentile(ally, 0.5), np.nanpercentile(ally, 99.5))
    for ax, f in zip(axes.flat, fields):
        t = load(f)
        nn = _hexbin_cmd(ax, t, b, r, y, xlim=xlim, ylim=ylim)
        ax.set_title(f"{TITLE[f]} ({f})   N={nn}", fontsize=11)
    fig.suptitle(f"{b.upper()} - {r.upper()}  vs  {y.upper()}   (Vega)",
                 fontsize=14)
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fn = f"{OUT}/compare_cmd_{tag}.png"
    fig.savefig(fn, dpi=130)
    plt.close(fig)
    print("wrote", fn)


def _panel_color_range(panels, pad=0.03):
    """Union (full inclusive) data range of both colors across all panels."""
    xs, ys = [], []
    for f, bx, rx, by, ry in panels:
        t = load(f)
        m = good(t, bx, rx, by, ry)
        cx = mag(t, bx)[m] - mag(t, rx)[m]
        cy = mag(t, by)[m] - mag(t, ry)[m]
        ok = np.isfinite(cx) & np.isfinite(cy)
        xs.append(cx[ok]); ys.append(cy[ok])
    xs = np.concatenate(xs); ys = np.concatenate(ys)
    def rng(a):
        lo, hi = np.nanmin(a), np.nanmax(a)
        d = (hi - lo) * pad
        return (lo - d, hi + d)
    return rng(xs), rng(ys)


def has(field, *filts):
    return all(f in FILTERS[field] for f in filts)


def stacked_cmd_grid(fields, cmds, tag, suptitle):
    """Rows = CMDs, columns = fields (side-by-side).  Each ROW shares its x/y
    limits across fields so the same CMD is directly comparable left-to-right.
    Fields lacking a row's filters get a blanked cell."""
    nrow, ncol = len(cmds), len(fields)
    fig, axes = plt.subplots(nrow, ncol, figsize=(4.6 * ncol, 4.4 * nrow),
                             squeeze=False)
    for i, (b, r, y) in enumerate(cmds):
        # shared row range from the fields that have the bands
        cc, yy = [], []
        for f in fields:
            if not has(f, b, r, y):
                continue
            t = load(f)
            m = good(t, b, r, y)
            cc.append(mag(t, b)[m] - mag(t, r)[m])
            yy.append(mag(t, y)[m])
        if cc:
            cc = np.concatenate(cc); yy = np.concatenate(yy)
            cc = cc[np.isfinite(cc)]; yy = yy[np.isfinite(yy)]
            xlim = (np.nanpercentile(cc, 0.5), np.nanpercentile(cc, 99.5))
            ylim = (np.nanpercentile(yy, 0.5), np.nanpercentile(yy, 99.5))
        else:
            xlim = ylim = None
        for j, f in enumerate(fields):
            ax = axes[i, j]
            if not has(f, b, r, y):
                ax.axis("off")
                ax.text(0.5, 0.5, f"{f}: no {b.upper()}/{r.upper()}/{y.upper()}",
                        ha="center", va="center", transform=ax.transAxes,
                        fontsize=9, color="0.5")
                continue
            t = load(f)
            nn = _hexbin_cmd(ax, t, b, r, y, xlim=xlim, ylim=ylim)
            ax.set_title(f"{TITLE[f]} ({f})  "
                         f"{b.upper()}-{r.upper()} vs {y.upper()}  N={nn}",
                         fontsize=9)
    fig.suptitle(suptitle, fontsize=14, y=0.997)
    fig.tight_layout(rect=[0, 0, 1, 1 - 0.32 / nrow])
    fn = f"{OUT}/cmdpanel_{tag}.png"
    fig.savefig(fn, dpi=120)
    plt.close(fig)
    print("wrote", fn)


def mixed_ccd_panels(panels, tag, suptitle, xlim=None, ylim=None):
    """Side-by-side CCDs where each field uses its OWN (nearest-wavelength) band
    quad -- for comparing a diagram that has no exact shared-filter analog.
    panels = [(field, bx, rx, by, ry), ...].  Pass xlim/ylim to SHARE axes across
    panels; pass 'inclusive' to share axes spanning ALL points; None autoscales
    each panel."""
    if xlim == "inclusive" or ylim == "inclusive":
        xr, yr = _panel_color_range(panels)
        xlim = xr if xlim == "inclusive" else xlim
        ylim = yr if ylim == "inclusive" else ylim
    fig, axes = plt.subplots(1, len(panels), figsize=(5.2 * len(panels), 5.4),
                             squeeze=False)
    for ax, (f, bx, rx, by, ry) in zip(axes.flat, panels):
        t = load(f)
        nn = _hexbin_ccd(ax, t, bx, rx, by, ry, xlim=xlim, ylim=ylim)
        ax.set_title(f"{TITLE[f]} ({f})   N={nn}\n"
                     f"{bx.upper()}-{rx.upper()} vs {by.upper()}-{ry.upper()}",
                     fontsize=10)
    fig.suptitle(suptitle, fontsize=14)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fn = f"{OUT}/compare_ccd_{tag}.png"
    fig.savefig(fn, dpi=130)
    plt.close(fig)
    print("wrote", fn)


def side_by_side_ccd(fields, bx, rx, by, ry, tag, xlim=None, ylim=None):
    fig, axes = plt.subplots(1, len(fields), figsize=(5.2 * len(fields), 5.4),
                             squeeze=False)
    if xlim is None or ylim is None:
        ax_, ay_ = [], []
        for f in fields:
            t = load(f)
            m = good(t, bx, rx, by, ry)
            ax_.append(mag(t, bx)[m] - mag(t, rx)[m])
            ay_.append(mag(t, by)[m] - mag(t, ry)[m])
        ax_ = np.concatenate(ax_); ay_ = np.concatenate(ay_)
        ax_ = ax_[np.isfinite(ax_)]; ay_ = ay_[np.isfinite(ay_)]
        if xlim is None:
            xlim = (np.nanpercentile(ax_, 1), np.nanpercentile(ax_, 99))
        if ylim is None:
            ylim = (np.nanpercentile(ay_, 1), np.nanpercentile(ay_, 99))
    for ax, f in zip(axes.flat, fields):
        t = load(f)
        nn = _hexbin_ccd(ax, t, bx, rx, by, ry, xlim=xlim, ylim=ylim)
        ax.set_title(f"{TITLE[f]} ({f})   N={nn}", fontsize=11)
    fig.suptitle(f"{bx.upper()}-{rx.upper()}  vs  {by.upper()}-{ry.upper()}"
                 f"   (Vega)", fontsize=14)
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fn = f"{OUT}/compare_ccd_{tag}.png"
    fig.savefig(fn, dpi=130)
    plt.close(fig)
    print("wrote", fn)


# ---------------------------------------------------------------------------
if __name__ == "__main__":
    warnings.simplefilter("ignore")

    # ===== per-field native CMD grids =====
    per_field_cmd_grid("wd1", [
        ("f115w", "f200w", "f200w"), ("f115w", "f212n", "f212n"),
        ("f200w", "f444w", "f444w"), ("f200w", "f405n", "f405n"),
        ("f187n", "f405n", "f405n"), ("f212n", "f323n", "f323n"),
        ("f277w", "f444w", "f444w"), ("f405n", "f466n", "f466n"),
    ])
    per_field_cmd_grid("wd2", [
        ("f115w", "f200w", "f200w"), ("f115w", "f212n", "f212n"),
        ("f182m", "f410m", "f410m"), ("f200w", "f444w", "f444w"),
        ("f187n", "f405n", "f405n"), ("f162m", "f212n", "f212n"),
        ("f335m", "f444w", "f444w"), ("f405n", "f466n", "f466n"),
    ])
    per_field_cmd_grid("w51", [
        ("f140m", "f210m", "f210m"), ("f140m", "f405n", "f405n"),
        ("f182m", "f410m", "f410m"), ("f162m", "f210m", "f210m"),
        ("f187n", "f405n", "f405n"), ("f210m", "f480m", "f480m"),
        ("f335m", "f480m", "f480m"), ("f360m", "f480m", "f480m"),
    ])

    # ===== per-field native CCD grids =====
    per_field_ccd_grid("wd1", [
        ("f115w", "f200w", "f200w", "f444w"),
        ("f115w", "f212n", "f212n", "f405n"),
        ("f187n", "f405n", "f200w", "f444w"),
        ("f200w", "f277w", "f277w", "f444w"),
        ("f212n", "f323n", "f405n", "f466n"),
        ("f150w", "f200w", "f200w", "f405n"),
    ])
    per_field_ccd_grid("wd2", [
        ("f115w", "f200w", "f200w", "f444w"),
        ("f162m", "f182m", "f335m", "f410m"),
        ("f187n", "f405n", "f200w", "f444w"),
        ("f182m", "f212n", "f335m", "f444w"),
        ("f212n", "f323n", "f405n", "f466n"),
        ("f150w", "f200w", "f300m", "f405n"),
    ])
    per_field_ccd_grid("w51", [
        ("f140m", "f210m", "f210m", "f405n"),
        ("f162m", "f182m", "f335m", "f410m"),
        ("f187n", "f405n", "f210m", "f480m"),
        ("f182m", "f210m", "f360m", "f480m"),
        ("f335m", "f410m", "f410m", "f480m"),
        ("f140m", "f162m", "f182m", "f210m"),
    ])

    # ===== side-by-side-by-side: common to ALL THREE (F187N, F405N) =====
    side_by_side_cmd(["wd1", "wd2", "w51"], "f187n", "f405n", "f405n",
                     "ALL3_f187n-f405n_vs_f405n")
    side_by_side_cmd(["wd1", "wd2", "w51"], "f187n", "f405n", "f187n",
                     "ALL3_f187n-f405n_vs_f187n")

    # ===== wd1 vs wd2 (11 shared filters) =====
    side_by_side_cmd(["wd1", "wd2"], "f115w", "f444w", "f444w",
                     "WD1WD2_f115w-f444w_vs_f444w")
    side_by_side_cmd(["wd1", "wd2"], "f200w", "f405n", "f405n",
                     "WD1WD2_f200w-f405n_vs_f405n")
    side_by_side_cmd(["wd1", "wd2"], "f115w", "f212n", "f212n",
                     "WD1WD2_f115w-f212n_vs_f212n")
    side_by_side_ccd(["wd1", "wd2"], "f115w", "f200w", "f200w", "f444w",
                     "WD1WD2_115-200_200-444")
    side_by_side_ccd(["wd1", "wd2"], "f187n", "f405n", "f212n", "f323n",
                     "WD1WD2_187-405_212-323")

    # ===== wd2 vs w51 (6 shared: f162m f182m f187n f335m f405n f410m) =====
    side_by_side_cmd(["wd2", "w51"], "f182m", "f410m", "f410m",
                     "WD2W51_f182m-f410m_vs_f410m")
    side_by_side_cmd(["wd2", "w51"], "f187n", "f405n", "f405n",
                     "WD2W51_f187n-f405n_vs_f405n")
    side_by_side_ccd(["wd2", "w51"], "f162m", "f182m", "f335m", "f410m",
                     "WD2W51_162-182_335-410")
    side_by_side_ccd(["wd2", "w51"], "f187n", "f405n", "f182m", "f410m",
                     "WD2W51_187-405_182-410")

    # ===== closest equivalent to w51's F162M-F210M vs F360M-F480M =====
    # w51 = exact; wd2/wd1 = nearest-wavelength substitutes:
    #   F162M -> F162M(wd2)/F164N(wd1) ; F210M -> F212N ;
    #   F360M -> F335M(wd2)/F323N(wd1) ; F480M -> F466N
    mixed_ccd_panels([
        ("wd1", "f164n", "f212n", "f323n", "f466n"),
        ("wd2", "f162m", "f212n", "f335m", "f466n"),
        ("w51", "f162m", "f210m", "f360m", "f480m"),
    ], "ALL3_shortcolor_vs_longcolor_nearest",
        "Short-wave color vs long-wave (3-5um) color   "
        "[w51: F162M-F210M vs F360M-F480M; wd1/wd2 nearest-band equivalents]",
        xlim=(-0.25, 4.5), ylim=(-1.0, 1.5))

    # a second long-color variant using F444W (broadband, better populated) for F480M
    mixed_ccd_panels([
        ("wd1", "f164n", "f212n", "f323n", "f444w"),
        ("wd2", "f162m", "f212n", "f335m", "f444w"),
        ("w51", "f162m", "f210m", "f360m", "f480m"),
    ], "ALL3_shortcolor_vs_longcolor_broadband",
        "Short vs long color (F444W long anchor for wd1/wd2)   "
        "[w51: F162M-F210M vs F360M-F480M]")

    # ===== extra band pairs =====
    side_by_side_cmd(["wd1", "wd2"], "f150w", "f444w", "f444w",
                     "WD1WD2_f150w-f444w_vs_f444w")
    side_by_side_ccd(["wd1", "wd2"], "f150w", "f200w", "f277w", "f444w",
                     "WD1WD2_150-200_277-444")
    side_by_side_ccd(["wd2", "w51"], "f162m", "f182m", "f187n", "f405n",
                     "WD2W51_162-182_187-405")

    # ======== ALL side-by-side CMDs we can build (stacked: rows=CMD, cols=field) ========
    # --- wd2 vs w51 (6 shared MB/NB: f162m f182m f187n f335m f405n f410m) ---
    #     includes the requested 162-182 and 335-410 colors.
    stacked_cmd_grid(["wd2", "w51"], [
        ("f162m", "f182m", "f182m"),
        ("f335m", "f410m", "f410m"),
        ("f162m", "f410m", "f410m"),
        ("f182m", "f410m", "f410m"),
        ("f162m", "f335m", "f335m"),
        ("f182m", "f335m", "f335m"),
        ("f187n", "f405n", "f405n"),
        ("f335m", "f405n", "f405n"),
    ], "WD2W51_shared", "wd2 vs w51 - all shared-filter CMDs (Vega)")

    # --- wd1 vs wd2 (11 shared broad/narrow bands) ---
    stacked_cmd_grid(["wd1", "wd2"], [
        ("f115w", "f200w", "f200w"),
        ("f115w", "f212n", "f212n"),
        ("f150w", "f200w", "f200w"),
        ("f200w", "f405n", "f405n"),
        ("f200w", "f444w", "f444w"),
        ("f115w", "f444w", "f444w"),
        ("f277w", "f444w", "f444w"),
        ("f212n", "f405n", "f405n"),
        ("f323n", "f466n", "f466n"),
        ("f405n", "f466n", "f466n"),
        ("f187n", "f405n", "f405n"),
    ], "WD1WD2_shared", "wd1 vs wd2 - all shared-filter CMDs (Vega)")

    # --- all three (only F187N, F405N common) ---
    stacked_cmd_grid(["wd1", "wd2", "w51"], [
        ("f187n", "f405n", "f405n"),
        ("f187n", "f405n", "f187n"),
    ], "ALL3_shared", "wd1 vs wd2 vs w51 - shared-filter CMDs (F187N/F405N only)")

    print("done")