#!/usr/bin/env python
"""Spatial source-density maps + radial profiles for wd1/wd2/w51.

Direct test of the claim that Wd1/Wd2's large catalogs are FIELD-dominated (line-
of-sight disk confusion), not cluster-dominated: if the counts were the cluster,
the radial density profile would fall steeply to ~0 off the core; if it's the
Galactic field, it plateaus on a high flat floor that the cluster only bumps.

Left: 2D source-density map (marks the density-peak center).
Right: azimuthally-averaged surface density vs radius, with the field floor
(median of the outer annuli) marked; the cluster is the excess above it.
"""
import os
import warnings
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

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"}


def field(f):
    t = Table.read(CATS[f])
    sc = SkyCoord(t["skycoord_ref"])
    return sc.ra.deg, sc.dec.deg


def analyze(f, ax_map, ax_rad):
    ra, dec = field(f)
    cosd = np.cos(np.radians(np.median(dec)))
    # density-peak center
    H, xe, ye = np.histogram2d(ra, dec, bins=100)
    i, j = np.unravel_index(np.argmax(H), H.shape)
    c_ra = 0.5 * (xe[i] + xe[i + 1])
    c_dec = 0.5 * (ye[j] + ye[j + 1])

    # ---- 2D density map (arcmin offsets from center) ----
    dx = (ra - c_ra) * cosd * 60.0
    dy = (dec - c_dec) * 60.0
    hb = ax_map.hexbin(dx, dy, gridsize=90, bins="log", cmap="magma", mincnt=1)
    ax_map.plot(0, 0, "c+", ms=14, mew=2)
    ax_map.set_aspect("equal")
    ax_map.set_xlabel("dRA (arcmin)")
    ax_map.set_ylabel("dDec (arcmin)")
    ax_map.set_title(f"{TITLE[f]} ({f}) - source density  (N={len(ra)})",
                     fontsize=10)
    plt.colorbar(hb, ax=ax_map, label="log N/cell", fraction=0.046)

    # ---- radial surface-density profile ----
    r = np.hypot(dx, dy)  # arcmin
    rmax = np.percentile(r, 98)
    edges = np.linspace(0, rmax, 26)
    cen = 0.5 * (edges[:-1] + edges[1:])
    area = np.pi * (edges[1:] ** 2 - edges[:-1] ** 2)  # arcmin^2
    cnt, _ = np.histogram(r, bins=edges)
    dens = cnt / area
    derr = np.sqrt(cnt) / area
    # field floor = median density of the outer third of annuli
    floor = np.median(dens[int(len(dens) * 0.66):])
    ax_rad.errorbar(cen, dens, yerr=derr, fmt="o-", ms=4, color="k",
                    label="observed")
    ax_rad.axhline(floor, color="tab:red", ls="--",
                   label=f"field floor ~{floor:.0f}/arcmin$^2$")
    # cluster excess: integrate (dens-floor)*area over annuli where >floor
    exc = np.clip(dens - floor, 0, None)
    n_cluster = np.sum(exc * area)
    n_total = len(ra)
    ax_rad.set_xlabel("radius from peak (arcmin)")
    ax_rad.set_ylabel("surface density (N/arcmin$^2$)")
    ax_rad.set_yscale("log")
    ax_rad.set_title(f"radial profile - cluster excess ~{n_cluster:.0f} "
                     f"({100*n_cluster/n_total:.0f}% of {n_total})", fontsize=10)
    ax_rad.legend(fontsize=8)
    return dict(field=f, n_total=n_total, floor=floor,
                n_cluster=n_cluster, frac_cluster=n_cluster / n_total)


def main():
    warnings.simplefilter("ignore")
    rows = []
    for f in ["wd1", "wd2", "w51"]:
        fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 5.3))
        rows.append(analyze(f, a1, a2))
        fig.tight_layout()
        fn = f"{OUT}/spatial_{f}.png"
        fig.savefig(fn, dpi=130)
        plt.close(fig)
        print("wrote", fn)
    print("\nfield: N_total  field_floor(/arcmin^2)  N_cluster_excess  cluster_frac")
    for r in rows:
        print(f"  {r['field']}: {r['n_total']:7d}  {r['floor']:8.0f}  "
              f"{r['n_cluster']:8.0f}  {100*r['frac_cluster']:.0f}%")


if __name__ == "__main__":
    main()