#!/usr/bin/env python
"""Cross-tool comparison of band-merged catalogs.

Produces:
  1. Pairwise positional matches between each tool and the iter3 reference
  2. Flux ratio histograms and scatter plots per filter
  3. Uncertainty distributions
  4. Recovery / new-source fractions
  5. Residual statistics in a spatial mask-defined crowded vs sparse subregion

Outputs written to benchmark/writeup/figures/ and benchmark/writeup/tables/.
"""
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, join
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"
TABDIR = config.BENCHMARK_ROOT / "writeup" / "tables"
FIGDIR.mkdir(parents=True, exist_ok=True)
TABDIR.mkdir(parents=True, exist_ok=True)

TOOLS = ["crowdsource", "starbug2", "dolphot", "photutils_iter3"]


def load_bandmerged(tool: str) -> Table:
    p = config.tool_bandmerged(tool)
    if not p.exists():
        print(f"WARN: {p} does not exist (tool {tool} not yet run)")
        return None
    return Table.read(p)


def positional_match(t_ref: Table, t_other: Table, radius_arcsec: float):
    c_ref = SkyCoord(t_ref["ra"] * u.deg, t_ref["dec"] * u.deg)
    c_oth = SkyCoord(t_other["ra"] * u.deg, t_other["dec"] * u.deg)
    idx, sep2d, _ = c_ref.match_to_catalog_sky(c_oth)
    match = sep2d.arcsec < radius_arcsec
    return idx, sep2d.arcsec, match


def compare_tool_vs_ref(tool: str, ref: Table, filters: list[str]) -> dict:
    t = load_bandmerged(tool)
    if t is None:
        return None
    idx, sep, match = positional_match(ref, t, config.MATCH_RADIUS_ARCSEC)
    results = {"tool": tool, "n_ref": len(ref), "n_tool": len(t), "n_matched": int(match.sum())}

    # Per-filter flux ratio
    for filt in filters:
        ref_flux = ref[f"{filt}_flux"]
        ref_err = ref[f"{filt}_flux_err"]
        tool_flux = t[f"{filt}_flux"][idx]
        tool_err = t[f"{filt}_flux_err"][idx]
        both = match & np.isfinite(ref_flux) & np.isfinite(tool_flux) & (ref_flux > 0) & (tool_flux > 0)
        ratio = tool_flux[both] / ref_flux[both]
        results[f"{filt}_n_both"] = int(both.sum())
        results[f"{filt}_median_ratio"] = float(np.median(ratio)) if both.sum() > 0 else np.nan
        results[f"{filt}_mad_ratio"] = float(np.median(np.abs(ratio - np.median(ratio)))) if both.sum() > 0 else np.nan

        # Plot flux-vs-flux
        fig, ax = plt.subplots(figsize=(5, 5))
        ax.loglog(ref_flux[both], tool_flux[both], ",", alpha=0.3)
        lo, hi = np.nanmin(ref_flux[both]), np.nanmax(ref_flux[both])
        ax.plot([lo, hi], [lo, hi], "r-", lw=0.8)
        ax.set_xlabel(f"iter3 flux ({filt})")
        ax.set_ylabel(f"{tool} flux ({filt})")
        ax.set_title(f"{tool} vs iter3, {filt}  n={both.sum()}")
        fig.tight_layout()
        fig.savefig(FIGDIR / f"flux_scatter_{tool}_vs_iter3_{filt}.png", dpi=120)
        plt.close(fig)

        # Uncertainty histogram
        fig, ax = plt.subplots(figsize=(5, 4))
        rel_err_ref = ref_err[both] / ref_flux[both]
        rel_err_tool = tool_err[both] / tool_flux[both]
        bins = np.logspace(-3, 0.5, 50)
        ax.hist(rel_err_ref[np.isfinite(rel_err_ref)], bins=bins, alpha=0.5, label="iter3")
        ax.hist(rel_err_tool[np.isfinite(rel_err_tool)], bins=bins, alpha=0.5, label=tool)
        ax.set_xscale("log")
        ax.set_xlabel("relative flux uncertainty")
        ax.set_ylabel("N")
        ax.set_title(f"uncertainty: {tool} vs iter3, {filt}")
        ax.legend()
        fig.tight_layout()
        fig.savefig(FIGDIR / f"relerr_{tool}_vs_iter3_{filt}.png", dpi=120)
        plt.close(fig)
    return results


def main():
    ref = load_bandmerged("photutils_iter3")
    if ref is None:
        print("Reference catalog (photutils_iter3) not available yet.")
        return
    rows = []
    for tool in ["crowdsource", "starbug2", "dolphot"]:
        row = compare_tool_vs_ref(tool, ref, config.FILTERS)
        if row is not None:
            rows.append(row)
    if rows:
        summary = Table(rows)
        summary.write(TABDIR / "comparison_summary.ecsv", overwrite=True)
        print(summary)


if __name__ == "__main__":
    main()