#!/usr/bin/env python
"""Difference-vs-flux plots: per filter, per alt-tool, vs photutils iter1.

Produces:
  1. one figure per filter, 3 subplots (one per alt tool):
       y = log10(alt_flux / iter1_flux), x = log10(iter1_flux)
     with binned median + 1.4826 * MAD shaded band.
  2. one combined "all filters" figure per alt tool (5 subplots).
  3. dedicated F470N "narrow-band-405-analog" figure with extra detail
     (linear delta-flux vs flux scatter + log-ratio panels).

iter1 is the primary reference (apples-to-apples per-frame blind detection).
A secondary set of figures uses iter3 as reference for context.
"""
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" / "fluxdiff"
FIGDIR.mkdir(parents=True, exist_ok=True)

ALT_TOOLS = ["crowdsource", "starbug2", "dolphot"]


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_bm(tool: str) -> Table | None:
    p = config.tool_bandmerged(tool)
    if not p.exists():
        return None
    return Table.read(p)


def matched_pair(ref: Table, alt: Table, filt: str, radius_arcsec: float):
    """Return (ref_flux[both], alt_flux[both], ref_err, alt_err) for matched sources."""
    c_r = SkyCoord(ref["ra"] * u.deg, ref["dec"] * u.deg)
    c_a = SkyCoord(alt["ra"] * u.deg, alt["dec"] * u.deg)
    idx, sep, _ = c_r.match_to_catalog_sky(c_a)
    matched = sep.arcsec < radius_arcsec

    rf = safe(ref, f"{filt}_flux")
    re = safe(ref, f"{filt}_flux_err")
    af_all = safe(alt, f"{filt}_flux")
    ae_all = safe(alt, f"{filt}_flux_err")
    af = af_all[idx]
    ae = ae_all[idx]
    both = matched & np.isfinite(rf) & np.isfinite(af) & (rf > 0) & (af > 0)
    return rf[both], af[both], re[both], ae[both]


def binned_stats(x: np.ndarray, y: np.ndarray, n_bins: int = 12, min_per_bin: int = 30):
    """Return bin centers and (median, mad) of y per equal-count bin in x."""
    if len(x) < min_per_bin * 2:
        return None
    order = np.argsort(x)
    x_s, y_s = x[order], y[order]
    edges = np.linspace(0, len(x_s), n_bins + 1).astype(int)
    centers = []
    medians = []
    mads = []
    for i in range(n_bins):
        s, e = edges[i], edges[i + 1]
        if e - s < min_per_bin:
            continue
        x_chunk = x_s[s:e]
        y_chunk = y_s[s:e]
        centers.append(float(np.median(x_chunk)))
        med = float(np.median(y_chunk))
        medians.append(med)
        mads.append(1.4826 * float(np.median(np.abs(y_chunk - med))))
    return np.array(centers), np.array(medians), np.array(mads)


def plot_per_filter_per_tool(ref_tool: str = "photutils_iter1"):
    """Per-filter figure with 3 alt-tool subplots."""
    ref = load_bm(ref_tool)
    if ref is None:
        print(f"missing {ref_tool}")
        return
    alts = {tool: load_bm(tool) for tool in ALT_TOOLS}
    tag = ref_tool.replace("photutils_", "")

    for filt in config.FILTERS:
        fig, axes = plt.subplots(1, 3, figsize=(14, 4.5), sharey=True)
        for ax, tool in zip(axes, ALT_TOOLS):
            alt = alts[tool]
            if alt is None:
                ax.set_title(f"{tool} (missing)")
                continue
            rf, af, _, _ = matched_pair(ref, alt, filt, config.MATCH_RADIUS_ARCSEC)
            if len(rf) == 0:
                ax.set_title(f"{tool} (no matches)")
                continue
            log_ratio = np.log10(af / rf)
            log_ratio -= np.median(log_ratio)
            log_x = np.log10(rf)
            ax.plot(log_x, log_ratio, ",", alpha=0.25, color="C0")
            stats = binned_stats(log_x, log_ratio)
            if stats is not None:
                xb, mb, sb = stats
                ax.plot(xb, mb, "-", color="C3", lw=1.6, label="binned median")
                ax.fill_between(xb, mb - sb, mb + sb, color="C3", alpha=0.2,
                                 label="±1.4826·MAD")
            ax.axhline(0, color="k", lw=0.5)
            ax.set_xlabel(f"log10({filt} ref flux)")
            ax.set_title(f"{tool}  N={len(rf)}")
            sigma = 1.4826 * np.median(np.abs(log_ratio))
            ax.text(0.02, 0.95, f"σ={sigma:.3f} dex",
                     transform=ax.transAxes, va="top", ha="left",
                     bbox=dict(boxstyle="round", fc="white", alpha=0.7))
            if ax is axes[0]:
                ax.set_ylabel(r"log10(alt/ref) - log10(median)")
                ax.legend(loc="lower left", fontsize=8)
            ax.set_ylim(-1.0, 1.0)
        fig.suptitle(f"{filt}: alt-tool flux vs {ref_tool} (matched within {config.MATCH_RADIUS_ARCSEC}\")")
        fig.tight_layout()
        outpath = FIGDIR / f"fluxdiff_per_tool_{filt}_vs_{tag}.png"
        fig.savefig(outpath, dpi=120)
        plt.close(fig)
        print(f"  wrote {outpath}")


def plot_all_filters_per_tool(ref_tool: str = "photutils_iter1"):
    """One figure per alt tool, with 5 filter subplots in a row."""
    ref = load_bm(ref_tool)
    if ref is None:
        return
    alts = {tool: load_bm(tool) for tool in ALT_TOOLS}
    tag = ref_tool.replace("photutils_", "")

    for tool in ALT_TOOLS:
        alt = alts[tool]
        if alt is None:
            continue
        fig, axes = plt.subplots(1, 5, figsize=(22, 4.5), sharey=True)
        for ax, filt in zip(axes, config.FILTERS):
            rf, af, _, _ = matched_pair(ref, alt, filt, config.MATCH_RADIUS_ARCSEC)
            if len(rf) == 0:
                ax.set_title(f"{filt} (no matches)"); continue
            log_ratio = np.log10(af / rf)
            log_ratio -= np.median(log_ratio)
            log_x = np.log10(rf)
            ax.plot(log_x, log_ratio, ",", alpha=0.25, color="C0")
            stats = binned_stats(log_x, log_ratio)
            if stats is not None:
                xb, mb, sb = stats
                ax.plot(xb, mb, "-", color="C3", lw=1.5)
                ax.fill_between(xb, mb - sb, mb + sb, color="C3", alpha=0.2)
            ax.axhline(0, color="k", lw=0.5)
            ax.set_xlabel(f"log10({filt} ref flux)")
            ax.set_title(f"{filt}  N={len(rf)}")
            sigma = 1.4826 * np.median(np.abs(log_ratio))
            ax.text(0.02, 0.95, f"σ={sigma:.3f} dex",
                     transform=ax.transAxes, va="top", ha="left",
                     bbox=dict(boxstyle="round", fc="white", alpha=0.7))
            if ax is axes[0]:
                ax.set_ylabel(r"log10(alt/ref) - log10(median)")
            ax.set_ylim(-1.0, 1.0)
        fig.suptitle(f"{tool}: flux ratio vs {ref_tool} for matched sources")
        fig.tight_layout()
        outpath = FIGDIR / f"fluxdiff_per_filter_{tool}_vs_{tag}.png"
        fig.savefig(outpath, dpi=120)
        plt.close(fig)
        print(f"  wrote {outpath}")


def plot_narrowband_focus(filt: str = "F470N", ref_tool: str = "photutils_iter1"):
    """Detailed figure for one filter (the F405N analog).

    Two-panel column per tool: top = log-ratio, bottom = absolute (alt - ref).
    """
    ref = load_bm(ref_tool)
    if ref is None:
        return
    alts = {tool: load_bm(tool) for tool in ALT_TOOLS}
    tag = ref_tool.replace("photutils_", "")

    fig, axes = plt.subplots(2, 3, figsize=(14, 8), sharex=True)
    for col, tool in enumerate(ALT_TOOLS):
        alt = alts[tool]
        if alt is None:
            continue
        rf, af, re_, ae_ = matched_pair(ref, alt, filt, config.MATCH_RADIUS_ARCSEC)
        if len(rf) == 0:
            continue
        log_x = np.log10(rf)

        # top: log-ratio, normalised
        log_ratio = np.log10(af / rf)
        log_ratio_norm = log_ratio - np.median(log_ratio)
        ax = axes[0, col]
        ax.plot(log_x, log_ratio_norm, ",", alpha=0.3, color="C0")
        stats = binned_stats(log_x, log_ratio_norm)
        if stats is not None:
            xb, mb, sb = stats
            ax.plot(xb, mb, "-", color="C3", lw=1.5)
            ax.fill_between(xb, mb - sb, mb + sb, color="C3", alpha=0.2)
        ax.axhline(0, color="k", lw=0.5)
        sigma = 1.4826 * np.median(np.abs(log_ratio_norm))
        ax.set_title(f"{tool}  N={len(rf)}  σ={sigma:.3f} dex")
        ax.set_ylabel(r"log10(alt/ref) - log10(median)")
        ax.set_ylim(-1.0, 1.0)

        # bottom: absolute delta normalised by ref's reported error
        # delta in units of σ_ref: (alt - ref) / sqrt(σ_alt² + σ_ref²)
        denom = np.sqrt(np.where(np.isfinite(re_), re_**2, 0.0)
                         + np.where(np.isfinite(ae_), ae_**2, 0.0))
        ok = denom > 0
        delta_sigma = (af[ok] - rf[ok]) / denom[ok]
        ax2 = axes[1, col]
        if ok.sum() > 0:
            ax2.plot(log_x[ok], delta_sigma, ",", alpha=0.3, color="C0")
            stats2 = binned_stats(log_x[ok], delta_sigma)
            if stats2 is not None:
                xb, mb, sb = stats2
                ax2.plot(xb, mb, "-", color="C3", lw=1.5)
                ax2.fill_between(xb, mb - sb, mb + sb, color="C3", alpha=0.2)
        ax2.axhline(0, color="k", lw=0.5)
        ax2.set_xlabel(f"log10({filt} ref flux)")
        ax2.set_ylabel(r"(alt - ref) / sqrt(σ_alt² + σ_ref²)")
        ax2.set_ylim(-10, 10)

    fig.suptitle(f"{filt} (narrow-band): alt-tool vs {ref_tool}")
    fig.tight_layout()
    outpath = FIGDIR / f"narrowband_focus_{filt}_vs_{tag}.png"
    fig.savefig(outpath, dpi=120)
    plt.close(fig)
    print(f"  wrote {outpath}")


def main():
    for ref in ["photutils_iter1", "photutils_iter3"]:
        print(f"=== reference: {ref} ===")
        plot_per_filter_per_tool(ref)
        plot_all_filters_per_tool(ref)
    # F470N is the narrow-band most analogous to F405N (we don't have F405N)
    plot_narrowband_focus("F470N", "photutils_iter1")
    plot_narrowband_focus("F470N", "photutils_iter3")


if __name__ == "__main__":
    main()