#!/usr/bin/env python
"""Color-magnitude diagrams per algorithm.

NOTE on filter substitution: this dataset doesn't have F356W or F444W.
Closest medium-band analogues are F335M (≈ F356W center) and F480M (≈ F444W
center). Likewise F210M substitutes for F200W. So:

  CMD axis  : F335M  vs  F335M − F480M
  blue cut  : F210M − F335M < threshold

Each tool has its own native flux unit, so we report instrumental magnitudes
(-2.5 log10 flux), then anchor each tool's per-filter zeropoint to the
photutils iter1 reference using the median of well-matched bright sources.
After the anchoring, the CMD shapes are directly comparable across tools.

Outputs (in benchmark/writeup/figures/v2/cmd/):
  cmd_all.png        : 2x3 grid, all matched stars
  cmd_blue.png       : 2x3 grid, blue-color-selected stars only
  cmd_individual_<tool>.png : single-panel version per tool
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord
from astropy.table import Table
from astropy import units as u

sys.path.insert(0, str(Path(__file__).resolve().parent))
import config  # noqa: E402

FIGDIR = config.BENCHMARK_ROOT / "writeup" / "figures" / "v2" / "cmd"
FIGDIR.mkdir(parents=True, exist_ok=True)

TOOLS = [
    ("photutils_iter1", "photutils iter1"),
    ("photutils_iter2", "photutils iter2"),
    ("photutils_iter3", "photutils iter3"),
    ("dolphot",         "DOLPHOT v3"),
    ("crowdsource",     "crowdsource"),
    ("starbug2",        "StarbugII"),
]

# CMD axis filters (substitutes for F356W / F444W / F200W on this dataset)
F_Y       = "F335M"  # ~F356W
F_RED     = "F480M"  # ~F444W
F_BLUE_CUT_BLUE = "F210M"  # ~F200W
F_BLUE_CUT_RED  = "F335M"  # the blue-cut color is F210M - F335M


def safe(t, col):
    if col not in t.colnames:
        return np.full(len(t), np.nan)
    return np.asarray(t[col]).astype(float)


def load(tool: str) -> Table | None:
    p = config.tool_bandmerged(tool)
    if not p.exists():
        print(f"missing {tool}")
        return None
    return Table.read(p)


def mag(flux: np.ndarray) -> np.ndarray:
    out = np.full_like(flux, np.nan, dtype=float)
    ok = np.isfinite(flux) & (flux > 0)
    out[ok] = -2.5 * np.log10(flux[ok])
    return out


def anchor_zeropoint(tool_tab: Table, ref_tab: Table, filt: str,
                     match_radius_arcsec: float = config.MATCH_RADIUS_ARCSEC) -> float:
    """Median (mag_ref - mag_tool) over matched bright sources for one filter.

    Returns a zeropoint such that `mag_tool + ZP ≈ mag_ref` for matched stars.
    """
    f_ref = safe(ref_tab, f"{filt}_flux")
    f_tool_all = safe(tool_tab, f"{filt}_flux")
    if not np.isfinite(f_ref).any() or not np.isfinite(f_tool_all).any():
        return 0.0
    # Match on sky
    c_r = SkyCoord(ref_tab["ra"] * u.deg, ref_tab["dec"] * u.deg)
    c_t = SkyCoord(tool_tab["ra"] * u.deg, tool_tab["dec"] * u.deg)
    idx, sep, _ = c_r.match_to_catalog_sky(c_t)
    matched = sep.arcsec < match_radius_arcsec
    f_tool = f_tool_all[idx]
    both = matched & np.isfinite(f_ref) & np.isfinite(f_tool) & (f_ref > 0) & (f_tool > 0)
    if both.sum() < 30:
        return 0.0
    # Use bright-half of the matched distribution to set the ZP
    m_ref = -2.5 * np.log10(f_ref[both])
    m_tool = -2.5 * np.log10(f_tool[both])
    bright = m_ref < np.percentile(m_ref, 30)  # top 30% brightest
    if bright.sum() < 10:
        bright = slice(None)
    return float(np.median(m_ref[bright] - m_tool[bright]))


def get_anchored(tool_tab: Table, ref_tab: Table | None, filt: str):
    """Return (mag_anchored, has_filter)."""
    flux = safe(tool_tab, f"{filt}_flux")
    m_raw = mag(flux)
    if ref_tab is not None and ref_tab is not tool_tab:
        zp = anchor_zeropoint(tool_tab, ref_tab, filt)
    else:
        zp = 0.0
    return m_raw + zp, np.any(np.isfinite(flux)), zp


def make_cmd_axes(ax, m_y, color, blue_mask, label, ylim, xlim, n_total):
    if blue_mask is None:
        ax.scatter(color, m_y, s=2, alpha=0.4, color="C0", linewidths=0)
        ax.set_title(f"{label}  N={int(np.sum(np.isfinite(m_y) & np.isfinite(color)))}")
    else:
        ok = np.isfinite(m_y) & np.isfinite(color)
        ax.scatter(color[ok & ~blue_mask], m_y[ok & ~blue_mask], s=2, alpha=0.15, color="0.6", linewidths=0)
        ax.scatter(color[ok & blue_mask], m_y[ok & blue_mask], s=3, alpha=0.7, color="C0", linewidths=0)
        ax.set_title(f"{label}  N(blue)={int(np.sum(ok & blue_mask))}/{int(np.sum(ok))}")
    ax.set_xlabel(f"{F_Y} - {F_RED}")
    ax.set_ylabel(F_Y + " (instrumental, anchored to iter1)")
    ax.set_ylim(ylim[1], ylim[0])  # mag axis inverted
    ax.set_xlim(xlim)
    ax.grid(True, alpha=0.3)


def main():
    cats = {tool: load(tool) for tool, _ in TOOLS}
    ref = cats["photutils_iter1"]
    if ref is None:
        print("ERR: photutils_iter1 not available")
        return

    # Build anchored mag(F335M), mag(F480M), mag(F210M) for each tool
    arr = {}
    for tool, label in TOOLS:
        t = cats[tool]
        if t is None:
            continue
        m_y, _, zpy = get_anchored(t, ref, F_Y)
        m_r, _, zpr = get_anchored(t, ref, F_RED)
        m_b, _, zpb = get_anchored(t, ref, F_BLUE_CUT_BLUE)
        arr[tool] = {
            "mag_y": m_y, "mag_r": m_r, "mag_b": m_b,
            "color_yr": m_y - m_r,
            "color_br": m_b - m_y,  # F210M - F335M
            "ra": np.asarray(t["ra"]), "dec": np.asarray(t["dec"]),
            "zp": (zpy, zpr, zpb),
        }
        print(f"{tool}: ZP({F_Y})={zpy:+.3f} ZP({F_RED})={zpr:+.3f} ZP({F_BLUE_CUT_BLUE})={zpb:+.3f}")

    # Determine common axis ranges from iter1
    m_y_ref = arr["photutils_iter1"]["mag_y"]
    color_ref = arr["photutils_iter1"]["color_yr"]
    color_b_ref = arr["photutils_iter1"]["color_br"]
    ok_ref = np.isfinite(m_y_ref) & np.isfinite(color_ref)
    ylim = (np.nanpercentile(m_y_ref[ok_ref], 0.5), np.nanpercentile(m_y_ref[ok_ref], 99.5))
    xlim = (np.nanpercentile(color_ref[ok_ref], 0.5), np.nanpercentile(color_ref[ok_ref], 99.5))
    # Pad a bit
    yspan = ylim[1] - ylim[0]
    xspan = xlim[1] - xlim[0]
    ylim = (ylim[0] - 0.1 * yspan, ylim[1] + 0.1 * yspan)
    xlim = (xlim[0] - 0.1 * xspan, xlim[1] + 0.1 * xspan)
    print(f"axis: F{F_Y} {ylim}; color {F_Y}-{F_RED} {xlim}")

    # Determine "blue" threshold from iter1's F210M-F335M distribution
    # "Blue" = bluer than the 25th percentile of finite values
    ok_b = np.isfinite(color_b_ref)
    blue_thresh = float(np.nanpercentile(color_b_ref[ok_b], 25.0))
    print(f"blue threshold ({F_BLUE_CUT_BLUE}-{F_BLUE_CUT_RED} <): {blue_thresh:.3f}")

    # Figure 1: all stars
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()
    for ax, (tool, label) in zip(axes, TOOLS):
        if tool not in arr:
            ax.set_axis_off(); continue
        a = arr[tool]
        make_cmd_axes(ax, a["mag_y"], a["color_yr"], None, label, ylim, xlim, len(a["mag_y"]))
    fig.suptitle(f"CMD: {F_Y} vs {F_Y}−{F_RED}  (all stars)\n"
                 f"NOTE: {F_Y} substitutes for F356W; {F_RED} substitutes for F444W on this dataset.",
                 fontsize=11)
    fig.tight_layout()
    fig.savefig(FIGDIR / "cmd_all.png", dpi=130)
    plt.close(fig)
    print(f"  wrote {FIGDIR/'cmd_all.png'}")

    # Figure 2: blue-color cut on F210M - F335M
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()
    for ax, (tool, label) in zip(axes, TOOLS):
        if tool not in arr:
            ax.set_axis_off(); continue
        a = arr[tool]
        blue_mask = np.isfinite(a["color_br"]) & (a["color_br"] < blue_thresh)
        make_cmd_axes(ax, a["mag_y"], a["color_yr"], blue_mask, label, ylim, xlim, len(a["mag_y"]))
    fig.suptitle(f"CMD: {F_Y} vs {F_Y}−{F_RED}  (blue: {F_BLUE_CUT_BLUE}−{F_BLUE_CUT_RED} < {blue_thresh:.2f}, "
                 f"iter1 25th percentile)\n"
                 f"NOTE: {F_BLUE_CUT_BLUE} substitutes for F200W; {F_Y} for F356W; {F_RED} for F444W.",
                 fontsize=11)
    fig.tight_layout()
    fig.savefig(FIGDIR / "cmd_blue.png", dpi=130)
    plt.close(fig)
    print(f"  wrote {FIGDIR/'cmd_blue.png'}")

    # Per-tool individual figures (higher resolution)
    for tool, label in TOOLS:
        if tool not in arr:
            continue
        a = arr[tool]
        fig, axes = plt.subplots(1, 2, figsize=(11, 6))
        # all
        make_cmd_axes(axes[0], a["mag_y"], a["color_yr"], None, f"{label}: all", ylim, xlim, 0)
        # blue
        blue_mask = np.isfinite(a["color_br"]) & (a["color_br"] < blue_thresh)
        make_cmd_axes(axes[1], a["mag_y"], a["color_yr"], blue_mask,
                      f"{label}: blue ({F_BLUE_CUT_BLUE}−{F_BLUE_CUT_RED}<{blue_thresh:.2f})", ylim, xlim, 0)
        fig.suptitle(f"{label}  CMD  {F_Y} vs {F_Y}−{F_RED}")
        fig.tight_layout()
        fig.savefig(FIGDIR / f"cmd_individual_{tool}.png", dpi=130)
        plt.close(fig)
    print(f"  wrote {FIGDIR}/cmd_individual_*.png")


if __name__ == "__main__":
    main()