#!/usr/bin/env python
"""Distance-corrected (absolute-magnitude) CMDs with reddening vectors, so the
three fields compare apples-to-apples despite different distances.

y-axis: M = mag_vega - DM   (distance modulus below).
reddening vector: A_lambda ~ lambda^-2.0 power law (Wang&Chen-like NIR slope),
normalized to A_F200W; arrow drawn for A_F200W = 3 mag.  This is an APPROXIMATE
law for the vector DIRECTION -- not a per-star de-reddening.

Distance moduli (kpc -> DM):
  wd1 4.05 -> 13.04   wd2 4.16 -> 13.10   w51 5.4 -> 13.66
"""
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"
CATS = {
    "wd1": f"{BASE}/w51-ewocs/catalogs/wd1_ewocs_merged_vega.fits",
    "wd2": 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"}
DM = {"wd1": 13.04, "wd2": 13.10, "w51": 13.66}

# effective wavelengths (um) from SVO; A ~ lambda^-2.0 normalized to F200W.
WL = {"f115w": 1.143, "f140m": 1.402, "f150w": 1.487, "f162m": 1.624,
      "f164n": 1.645, "f182m": 1.839, "f187n": 1.874, "f200w": 1.968,
      "f210m": 2.091, "f212n": 2.121, "f250m": 2.501, "f277w": 2.728,
      "f300m": 2.982, "f323n": 3.237, "f335m": 3.354, "f360m": 3.615,
      "f405n": 4.052, "f410m": 4.072, "f444w": 4.35, "f466n": 4.654,
      "f480m": 4.814}
AEXT = {f: (WL[f] / WL["f200w"]) ** (-2.0) for f in WL}  # A_f / A_F200W
QFIT_MAX, EMAG_MAX = 0.4, 0.15


def mag(t, fl):
    return np.asarray(t[f"mag_vega_{fl}"], 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(t[c], bool) if t[c].dtype == bool
            else np.asarray(np.nan_to_num(np.asarray(t[c], float)), bool))


def good(t, *filts):
    m = np.ones(len(t), bool)
    for fl in filts:
        mg = mag(t, fl)
        em = np.asarray(t[f"emag_ab_{fl}"], float)
        qf = np.asarray(t[f"qfit_{fl}"], 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 redvec(ax, b, r, y, av200=3.0, x0=None, y0=None):
    """Draw a reddening arrow for A_F200W=av200 in a (b-r) vs M_y CMD."""
    dcol = (AEXT[b] - AEXT[r]) * av200
    dmag = AEXT[y] * av200
    xl, yl = ax.get_xlim(), ax.get_ylim()
    if x0 is None:
        x0 = xl[0] + 0.12 * (xl[1] - xl[0])
    if y0 is None:
        y0 = yl[0] + 0.15 * (yl[1] - yl[0])  # yl already inverted
    ax.annotate("", xy=(x0 + dcol, y0 + dmag), xytext=(x0, y0),
                arrowprops=dict(arrowstyle="-|>", color="red", lw=2))
    ax.text(x0, y0 - 0.03 * (yl[0] - yl[1]),
            f"$A_{{F200W}}$={av200:g}", color="red", fontsize=8)


def side_by_side_absmag(fields, b, r, y, tag, av200=3.0):
    # shared axes across fields
    cc, yy = [], []
    for f in fields:
        t = Table.read(CATS[f])
        m = good(t, b, r, y)
        cc.append(mag(t, b)[m] - mag(t, r)[m])
        yy.append(mag(t, y)[m] - DM[f])
    cc = np.concatenate(cc); yy = np.concatenate(yy)
    cc = cc[np.isfinite(cc)]; yy = yy[np.isfinite(yy)]
    xlim = (np.nanpercentile(cc, 1), np.nanpercentile(cc, 99))
    ylim = (np.nanpercentile(yy, 0.5), np.nanpercentile(yy, 99.5))
    fig, axes = plt.subplots(1, len(fields), figsize=(5.2 * len(fields), 6),
                             squeeze=False)
    for ax, f in zip(axes.flat, fields):
        t = Table.read(CATS[f])
        m = good(t, b, r, y)
        color = mag(t, b)[m] - mag(t, r)[m]
        M = mag(t, y)[m] - DM[f]
        ok = np.isfinite(color) & np.isfinite(M)
        color, M = color[ok], M[ok]
        ext = (xlim[0], xlim[1], ylim[0], ylim[1])
        if len(color) > 400:
            ax.hexbin(color, M, gridsize=110, bins="log", cmap="viridis",
                      mincnt=1, extent=ext)
        else:
            ax.scatter(color, M, s=4, c="k", alpha=0.5, edgecolors="none")
        ax.set_xlim(xlim)
        ax.set_ylim(ylim[1], ylim[0])
        redvec(ax, b, r, y, av200=av200)
        ax.set_xlabel(f"{b.upper()} - {r.upper()}  (Vega)")
        ax.set_ylabel(f"$M_{{{y.upper()}}}$  (Vega, DM={DM[f]})")
        ax.set_title(f"{TITLE[f]} ({f})   N={len(color)}", fontsize=11)
    fig.suptitle(f"Distance-corrected CMD:  {b.upper()}-{r.upper()} vs "
                 f"M_{y.upper()}   (red arrow = reddening)", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fn = f"{OUT}/absmag_cmd_{tag}.png"
    fig.savefig(fn, dpi=130)
    plt.close(fig)
    print("wrote", fn)


def main():
    warnings.simplefilter("ignore")
    # all three (only F187N,F405N common)
    side_by_side_absmag(["wd1", "wd2", "w51"], "f187n", "f405n", "f405n",
                        "ALL3_187-405")
    # wd2 vs w51 shared medium bands
    side_by_side_absmag(["wd2", "w51"], "f182m", "f410m", "f410m",
                        "WD2W51_182-410")
    # wd1 vs wd2 broadbands
    side_by_side_absmag(["wd1", "wd2"], "f115w", "f200w", "f200w",
                        "WD1WD2_115-200")
    side_by_side_absmag(["wd1", "wd2"], "f200w", "f444w", "f444w",
                        "WD1WD2_200-444")


if __name__ == "__main__":
    main()