#!/usr/bin/env python3
"""Figures for the fp_kiii study (static PNGs)."""
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker  # noqa: F401
import numpy as np
import pandas as pd
from scipy import sparse
from scipy.spatial import cKDTree

import fpk_common as F

# palette (validated categorical order: blue, orange, aqua; neutral inks)
BLUE, ORANGE, AQUA, YELLOW = "#2a78d6", "#eb6834", "#1baf7a", "#eda100"
INK, INK2, GRID, SURF = "#0b0b0b", "#52514e", "#dddcd6", "#fcfcfb"
plt.rcParams.update({
    "figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF,
    "axes.edgecolor": INK2, "axes.labelcolor": INK, "xtick.color": INK2, "ytick.color": INK2,
    "text.color": INK, "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6,
    "axes.spines.top": False, "axes.spines.right": False, "font.size": 9.5,
    "axes.titlesize": 10.5, "axes.titleweight": "bold", "legend.frameon": False, "lines.linewidth": 2.0,
})


def fig_residuals():
    rs = pd.read_csv(F.resid_path("sdss"))
    r6 = pd.read_csv(os.path.join(F.DATA, "resid_6df.csv"))
    rh = pd.read_csv(os.path.join(F.DATA, "resid_howlett.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.0), sharey=True)
    bins = np.arange(-1.3, 0.9001, 0.05)
    xc = 0.5 * (bins[1:] + bins[:-1])
    panels = [
        (axes[0], rs.dI.values, "SDSS r band, own sample (N=%d)" % len(rs),
         rh.dI.values, "published, FP-clipped catalogue (left: Howlett+22; right: 6dFGSv with chi2<12)"),
        (axes[1], r6.dI.values, "6dFGSv J band (N=%d)" % len(r6),
         r6.dI.values[r6.in_6dFGSv_J.values == 1], None),
    ]
    for ax, d, title, dc, clab in panels:
        h = np.histogram(d, bins)[0] / len(d)
        hc = np.histogram(dc, bins)[0] / len(dc)
        sd = F.mad_sigma(d)
        gauss = np.diff(bins) * np.exp(-0.5 * ((xc - np.median(d)) / sd) ** 2) / (sd * np.sqrt(2 * np.pi))
        l1, = ax.plot(xc, gauss, color=INK2, lw=1.0, ls=":")
        l2, = ax.step(xc, np.where(hc > 0, hc, np.nan), where="mid", color=ORANGE, lw=1.6)
        l3, = ax.step(xc, np.where(h > 0, h, np.nan), where="mid", color=BLUE, lw=2.0)
        for f in (0.5, 0.75, 0.9):
            x = np.log10(1 - f)
            ax.axvline(x, color=INK2, lw=0.8, ls="--")
            ax.text(x - 0.012, 0.35, "f = %g" % f, rotation=90, ha="right", va="top", fontsize=8.5,
                    color=INK2)
        ax.set_yscale("log")
        ax.set_ylim(0.5 / len(d) if ax is axes[1] else 0.5 / len(rs), 0.5)
        ax.set_xlim(-1.3, 0.9)
        ax.set_title("%s; core width (MAD) %.3f dex" % (title, sd), loc="left")
        ax.set_xlabel(r"FP surface-brightness residual  $\delta_I$ = log $I_e$ - plane($\sigma$, $R_e$, z)  [dex]")
    axes[0].set_ylim(0.5 / len(rs), 0.5)
    axes[0].set_ylabel("fraction of sample per 0.05 dex")
    fig.legend([l3, l2, l1], ["sample without FP-outlier clipping (this work)",
                              "published FP-clipped catalogue (left: Howlett+22; right: 6dFGSv, chi2 < 12)",
                              "Gaussian with the core width"], loc="lower center", ncol=3, fontsize=8.5)
    fig.suptitle("Under-luminous tail of the Fundamental Plane. Dashed lines: where galaxies with a fraction f of "
                 "their starlight removed would be centred", fontsize=10.5, x=0.01, ha="left")
    fig.tight_layout(rect=(0, 0.06, 1, 0.95))
    fig.savefig(os.path.join(F.FIGS, "fig1_fp_residuals.png"), dpi=150)
    plt.close(fig)


def fig_limits():
    lim = pd.read_csv(os.path.join(F.DATA, "limits_individual_best.csv"))
    fin = pd.read_csv(os.path.join(F.DATA, "limits_final.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.3), gridspec_kw={"width_ratios": [1.45, 1]})
    ax = axes[0]
    series = [("SDSS-MGS own (unclipped)", "FP", BLUE, "SDSS r band, Fundamental Plane (N=%d)"),
              ("6dFGSv-J (unclipped)", "FP", ORANGE, "6dFGSv J band, Fundamental Plane (N=%d)"),
              ("SDSS-MGS own (unclipped)", "FJ", AQUA, "SDSS, Faber-Jackson cross-check (no radius used)")]
    handles = []
    for samp, rel, col, lab in series:
        s = lim[(lim["sample"] == samp) & (lim.relation == rel) & (lim.variant == "default")].sort_values("f")
        s = s[np.isfinite(s.q95_best) & (s.q95_best < 1)]
        h, = ax.plot(s.f, s.q95_best, color=col, marker="o", ms=5, label=lab % s.N.iloc[0] if "%d" in lab else lab)
        handles.append(h)
    red_fn = os.path.join(F.DATA, "limits_redext.csv")
    if os.path.exists(red_fn):
        rx = pd.read_csv(red_fn)
        rx = rx[(rx.relation == "FP") & np.isfinite(rx.q95_best) & (rx.q95_best < 1)].sort_values("f")
        h, = ax.plot(rx.f, rx.q95_best, color=YELLOW, marker="^", ms=4.5, lw=1.3,
                     label="SDSS FP, sample without the red colour edge (N=%d): also covers dimming that reddens"
                           % rx.N.iloc[0])
        handles.append(h)
    v = fin[(fin["sample"] == "SDSS-MGS own (unclipped)") & (fin.kind == "vetted") & (fin.q95 < 1)].sort_values("f")
    h, = ax.plot(v.f, v.q95, color=BLUE, ls="--", marker="s", ms=5, mfc=SURF, lw=1.4,
                 label="SDSS FP, counting only outliers left unexplained by image vetting")
    handles.append(h)
    # literature points (all quoted for f = 0.75)
    h1, = ax.plot([0.75], [0.028], marker="v", color=INK, ms=8, ls="none",
                  label="Annis 1999: 0 of 106 ellipticals (3/106; no selection correction)")
    h2, = ax.plot([0.757, 0.757], [0.03, 0.003], marker="D", color=INK, ms=6, ls="none", mfc=SURF,
                  label="Zackrisson+15, disks via Tully-Fisher: 3% conservative, 0.3% 'tentative'")
    handles += [h1, h2]
    ax.set_yscale("log")
    ax.set_ylim(5e-4, 1.5)
    ax.set_xlim(0.26, 0.99)
    ax.set_xlabel("fraction f of r-band (J-band for 6dFGSv) light removed, grey and spatially uniform")
    ax.set_ylabel("95% upper limit on the fraction of early-type galaxies")
    ax.set_title("Individual galaxies: upper limits for a single value of f", loc="left")
    ax = axes[1]
    raw = pd.read_csv(os.path.join(F.DATA, "limits_individual.csv"))
    for samp, col, lab in (("SDSS-MGS own (unclipped)", BLUE, "SDSS"), ("6dFGSv-J (unclipped)", ORANGE, "6dFGSv")):
        s = raw[(raw["sample"] == samp) & (raw.relation == "FP") & (raw.variant == "default") & (raw.kappa == 0)]
        s = s.sort_values("f")
        ax.plot(s.f, np.maximum(s.N_eff / s.N, 1e-5), color=col, marker="o", ms=5)
        s2 = raw[(raw["sample"] == samp) & (raw.relation == "FP") & (raw.variant == "snr_linear") & (raw.kappa == 0)]
        s2 = s2.sort_values("f")
        ax.plot(s2.f, np.maximum(s2.N_eff / s2.N, 1e-5), color=col, ls=":", lw=1.4)
    hp, = ax.plot([], [], color=INK2, ls=":", lw=1.4, label="same, if spectral S/N scaled linearly with flux (pessimistic)")
    handles.append(hp)
    ax.set_yscale("log")
    ax.set_ylim(1e-4, 1.5)
    ax.set_xlim(0.26, 0.99)
    ax.set_xlabel("fraction f of starlight removed")
    ax.set_ylabel(r"$N_{\rm eff}/N$")
    ax.set_title("Fraction of the sample that would survive the dimming", loc="left")
    fig.legend(handles=handles, loc="lower center", ncol=2, fontsize=8.0)
    fig.tight_layout(rect=(0, 0.2, 1, 1))
    fig.savefig(os.path.join(F.FIGS, "fig2_limits_individual.png"), dpi=150)
    plt.close(fig)


def fig_regions():
    rows = pd.read_csv(os.path.join(F.DATA, "limits_regions.csv"))
    fig, axes = plt.subplots(1, 4, figsize=(16.5, 5.8), gridspec_kw={"width_ratios": [1.15, 1, 1, 1]})
    # (a) observed vs shuffled distribution of S for SDSS, R = 10 Mpc
    ax = axes[0]
    rs = pd.read_csv(F.resid_path("sdss"))
    zt = pd.read_csv(os.path.join(F.DATA, "zeta_sdss.csv"))
    rs = rs.merge(zt, on="id")
    xyz = F.radec_to_xyz(rs.ra.values, rs["dec"].values, F.comoving_distance(rs.zdist.values))
    tree = cKDTree(xyz)
    pairs = tree.query_pairs(10.0, output_type="ndarray")
    n = len(rs)
    i = np.concatenate([pairs[:, 0], pairs[:, 1], np.arange(n)])
    j = np.concatenate([pairs[:, 1], pairs[:, 0], np.arange(n)])
    Aij = sparse.csr_matrix((np.ones(len(i), np.float32), (i, j)), shape=(n, n))
    N = np.asarray(Aij.sum(axis=1)).ravel()
    good = N >= 3
    S = (Aij @ rs.zeta.values)[good] / np.sqrt(N[good])
    rng = np.random.default_rng(1)
    bins = np.arange(-9.5, 8.01, 0.25)
    xc = 0.5 * (bins[1:] + bins[:-1])
    hs = np.zeros(len(xc))
    for _ in range(20):
        Ss = (Aij @ rng.permutation(rs.zeta.values))[good] / np.sqrt(N[good])
        hs += np.histogram(Ss, bins)[0] / 20.0
    h = np.histogram(S, bins)[0]
    ax.step(xc, np.where(hs > 0, hs, np.nan), where="mid", color=INK2, lw=1.4,
            label="residuals shuffled among galaxies (mean of 20)")
    ax.step(xc, np.where(h > 0, h, np.nan), where="mid", color=BLUE, lw=2.0, label="observed")
    ax.set_yscale("log")
    ax.set_xlabel(r"sphere statistic $S=\sum\zeta/\sqrt{N}$  (R = 10 Mpc)")
    ax.set_ylabel("number of spheres (centred on galaxies)")
    ax.set_title("SDSS: coherence of FP residuals", loc="left")
    ax.text(0.98, 0.97, "dim tail (S < -5):\nrich clusters\n(Coma, A2593,\nA2142, ...)", transform=ax.transAxes,
            fontsize=8.5, color=INK2, va="top", ha="right")
    ax.set_xlim(-9.5, 8)
    ax.legend(loc="upper left", bbox_to_anchor=(0.0, -0.17), fontsize=8.5, ncol=1)
    # (b, c) filling-fraction limits
    for ax, samp, title in ((axes[1], "sdss", "SDSS (z = 0.02-0.12)"), (axes[2], "6df", "6dFGSv (cz = 3000-16120 km/s)")):
        for f, col in ((0.3, BLUE), (0.5, ORANGE), (0.75, AQUA)):
            s = rows[(rows["sample"] == samp) & (rows["mode"] == "volume") & (rows.variant == "raw") & (rows.f == f)]
            s = s.sort_values("R_mpc")
            ax.plot(s.R_mpc, s.fill95, color=col, marker="o", ms=5, label="dimmed by f = %.2f" % f)
            ax.plot(s.R_mpc, s.fill95_fix, color=col, ls=":", lw=1.4, marker="s", ms=3.5)
        ax.plot([], [], color=INK2, marker="o", ms=5, label="nothing beyond the most extreme observed sphere")
        ax.plot([], [], color=INK2, ls=":", lw=1.4, marker="s", ms=3.5,
                label="fixed threshold S < -5, all observed spheres\nbelow it counted as candidates")
        ax.set_xscale("log")
        ax.set_yscale("log")
        ax.set_xticks([5, 10, 20])
        ax.set_xticklabels(["5", "10", "20"])
        ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
        ax.set_xlim(4, 25)
        ax.set_ylim(1e-3, 1.0)
        ax.set_xlabel("region radius R [Mpc]")
        ax.set_title(title, loc="left")
    axes[1].legend(loc="upper left", bbox_to_anchor=(-0.1, -0.17), fontsize=8.5, ncol=2)
    axes[1].set_ylabel("95% upper limit on volume-filling fraction\n(assumes regions placed uniformly in volume)")
    # (d) how much of the survey volume the limit was actually derived from
    ax = axes[3]
    fin = pd.read_csv(os.path.join(F.DATA, "limits_regions_final.csv"))
    for samp, ls, mk in (("sdss", "-", "o"), ("6df", "--", "s")):
        for f, col in ((0.3, BLUE), (0.5, ORANGE), (0.75, AQUA)):
            g = fin[(fin["sample"] == samp) & (fin.f == f)].sort_values("R_mpc")
            ax.plot(g.R_mpc, g.tested_share_of_survey_volume, color=col, ls=ls, marker=mk, ms=4.5, lw=1.6)
    ax.plot([], [], color=INK2, ls="-", marker="o", ms=4.5, lw=1.6, label="SDSS")
    ax.plot([], [], color=INK2, ls="--", marker="s", ms=4.5, lw=1.6, label="6dFGSv")
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xticks([5, 10, 20])
    ax.set_xticklabels(["5", "10", "20"])
    ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    ax.set_xlim(4, 25)
    ax.set_ylim(2e-4, 1.0)
    ax.set_xlabel("region radius R [Mpc]")
    ax.set_ylabel("share of the survey volume in which a dimmed\nsphere would have been detected")
    ax.set_title("Volume actually tested", loc="left")
    ax.legend(loc="upper left", bbox_to_anchor=(0.0, -0.17), fontsize=8.5, ncol=2)
    fig.suptitle("Coherently dimmed regions: observed sphere statistics, limits on the volume-filling fraction "
                 "of dimmed spheres, and the share of the survey volume those limits were derived from",
                 fontsize=10.5, x=0.01, ha="left")
    fig.subplots_adjust(left=0.05, right=0.99, top=0.86, bottom=0.34, wspace=0.33)
    fig.savefig(os.path.join(F.FIGS, "fig3_regions.png"), dpi=150)
    plt.close(fig)


def fig_map():
    """RA - redshift map of the locally averaged, selection-corrected FP residual (spheres of 10 Mpc)."""
    pt = pd.read_csv(os.path.join(F.DATA, "region_extreme_patches.csv"))
    summ = {(s["sample"], s["detrend"]): s for s in json.load(open(os.path.join(F.DATA, "region_summary.json")))}
    fig, axes = plt.subplots(2, 1, figsize=(11.5, 8.6))
    for ax, samp, title in ((axes[0], "sdss", "SDSS own sample"), (axes[1], "6df", "6dFGSv")):
        d = pd.read_csv(F.resid_path(samp)).merge(pd.read_csv(os.path.join(F.DATA, "zeta_%s.csv" % samp)), on="id")
        sd = summ[(samp, False)]["sd_dex"]
        xyz = F.radec_to_xyz(d.ra.values, d["dec"].values, F.comoving_distance(d.zdist.values))
        tree = cKDTree(xyz)
        pairs = tree.query_pairs(10.0, output_type="ndarray")
        n = len(d)
        i = np.concatenate([pairs[:, 0], pairs[:, 1], np.arange(n)])
        j = np.concatenate([pairs[:, 1], pairs[:, 0], np.arange(n)])
        Aij = sparse.csr_matrix((np.ones(len(i), np.float32), (i, j)), shape=(n, n))
        N = np.asarray(Aij.sum(axis=1)).ravel()
        mean = (Aij @ d.zeta.values) / N * sd
        m = N >= 5
        order = np.argsort(np.abs(mean[m]))
        sc = ax.scatter(d.ra.values[m][order], d.zdist.values[m][order], c=mean[m][order], s=3, cmap="RdBu",
                        vmin=-0.12, vmax=0.12, linewidths=0, rasterized=True)
        p = pt[(pt["sample"] == samp) & (pt.variant == "raw") & (pt.R_mpc == 10) & pt.kind.str.contains("negative")]
        ax.scatter(p.ra, p.z, s=140, facecolors="none", edgecolors=INK, linewidths=1.2,
                   label="the 5 most significant dim spheres (R = 10 Mpc), numbered by significance")
        for k, (_, r) in enumerate(p.iterrows()):
            ax.annotate("%d" % (k + 1), (r.ra, r.z), xytext=(7, 5), textcoords="offset points", fontsize=8.5)
        ax.set_xlabel("right ascension [deg]")
        ax.set_ylabel("redshift")
        ax.set_title("%s: mean FP residual of the galaxies within 10 Mpc (spheres with N >= 5)" % title, loc="left")
        ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.13), fontsize=8.5)
        cb = fig.colorbar(sc, ax=ax, pad=0.01, fraction=0.03)
        cb.set_label("mean residual [dex]  (red = dimmer than the FP)")
        ax.grid(False)
    fig.tight_layout()
    fig.savefig(os.path.join(F.FIGS, "fig4_residual_map.png"), dpi=140)
    plt.close(fig)


def fig_environment():
    """Mean FP residual versus local density, and the anatomy of the most significant dim sphere."""
    env = pd.read_csv(os.path.join(F.DATA, "env_trend.csv"))
    fog = pd.read_csv(os.path.join(F.DATA, "region_patches_groupz.csv"))
    pt = pd.read_csv(os.path.join(F.DATA, "region_extreme_patches.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(12.5, 5.0), gridspec_kw={"width_ratios": [1.1, 1]})
    ax = axes[0]
    for (lab, g), col, mk in zip(env.groupby("sample", sort=False), (BLUE, ORANGE, AQUA), ("o", "s", "D")):
        x = np.arange(len(g))
        ax.errorbar(x, g.mean_dI_dex, yerr=g.sem_dex, color=col, marker=mk, ms=5, lw=1.6, capsize=2, label=lab)
    labs = ["<0.5", "0.5-1", "1-2", "2-4", "4-8", ">8"]
    ax.set_xticks(range(len(labs)))
    ax.set_xticklabels(labs)
    ax.axhline(0, color=INK2, lw=0.8)
    ax.set_xlabel("local over-density of sample galaxies (within 5 Mpc, redshift space; 1 = mean)")
    ax.set_ylabel(r"mean FP residual $\delta_I$ [dex]")
    ax.set_title("Galaxies in dense regions lie below the plane", loc="left")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.17), fontsize=8.5, ncol=3)
    # (b) decomposition for the top R = 10 Mpc sphere, compared with the like-for-like (group-redshift) trend
    ax = axes[1]
    dec = pd.read_csv(os.path.join(F.DATA, "coma_decomposition.csv")).iloc[0]
    cl = pd.read_csv(os.path.join(F.DATA, "cluster_spheres.csv"))
    cat = cl[(cl.cluster == "ACO 1656") & (cl.R_mpc == 10)].iloc[0]
    vals = [dec.mean_dI, dec.fog_part, dec.phot_part, dec.remainder, cat.mean_dI]
    errs = [dec.sem_dI, 0, dec.phot_err, dec.remainder_err, cat.sem_dI]
    names = ["scan-selected\nsphere\n(N=%d)" % dec.N, "finger-of-god\n(individual vs\ngroup redshift)",
             "r-band\nphotometry\n(r - J residual)", "remainder", "sphere at the\ncatalogue centre\n(N=%d)" % cat.N]
    cols = [INK2, BLUE, BLUE, ORANGE, AQUA]
    ax.bar(range(5), vals, color=cols, width=0.6, yerr=errs, ecolor=INK, capsize=2, error_kw=dict(lw=0.9))
    for k, v in enumerate(vals):
        ax.text(k, v - errs[k] - 0.004, "%+.3f" % v, ha="center", va="top", fontsize=9, color=INK)
    ax.axhspan(dec.env_trend_group_z - dec.env_trend_group_z_err, dec.env_trend_group_z + dec.env_trend_group_z_err,
               color=ORANGE, alpha=0.18, lw=0)
    hl = ax.axhline(dec.env_trend_group_z, color=ORANGE, lw=1.0, ls="--",
                    label="mean residual of Howlett+22 galaxies (group redshifts) at over-density > 8, +- 1 sigma")
    ax.legend(handles=[hl], loc="upper center", bbox_to_anchor=(0.5, -0.2), fontsize=8.3)
    ax.set_xticks(range(5))
    ax.set_xticklabels(names, fontsize=8.3)
    ax.axhline(0, color=INK2, lw=0.8)
    ax.set_ylim(min(vals) - 0.03, 0.012)
    ax.set_ylabel("contribution to the mean residual [dex]")
    ax.set_title("Coma cluster, R = 10 Mpc: anatomy of the most\nsignificant dim sphere", loc="left")
    ax.grid(axis="x", visible=False)
    fig.tight_layout()
    fig.savefig(os.path.join(F.FIGS, "fig5_environment.png"), dpi=150)
    plt.close(fig)


def fig_redside():
    """Galaxies excluded by the red edge of the colour window: colour excess versus FP residual."""
    o = pd.read_csv(os.path.join(F.DATA, "redside_objects.csv"))
    fig, ax = plt.subplots(figsize=(7.4, 4.8))
    x = np.clip(o.excess_red.values, 0, 1.0)
    ax.axhspan(-1.3, -0.30, xmin=0.2 / 1.05, color=ORANGE, alpha=0.13, lw=0)
    ax.scatter(x, o.dI, s=7, color=BLUE, alpha=0.55, linewidths=0)
    for f in (0.5, 0.75):
        ax.axhline(np.log10(1 - f), color=INK2, lw=0.8, ls="--")
        ax.text(1.04, np.log10(1 - f) + 0.01, "f = %g" % f, ha="right", va="bottom", fontsize=8.5, color=INK2)
    ax.axvline(0.2, color=INK2, lw=0.8, ls=":")
    ax.text(0.62, -1.02, "where an old population dimmed by >= 50%\nby cloaking its fainter stars would lie\n"
            "(redder by >= 0.2 mag; Lacki 2019)", fontsize=8.5, color=INK2, ha="center")
    ax.set_xlim(0, 1.05)
    ax.set_ylim(-1.3, 0.6)
    ax.set_xlabel("g - r excess beyond the red edge of the colour window [mag] (values > 1 shown at 1)")
    ax.set_ylabel(r"FP surface-brightness residual $\delta_I$ [dex]")
    ax.set_title("The %d galaxies excluded only by the red colour edge" % len(o), loc="left")
    fig.tight_layout()
    fig.savefig(os.path.join(F.FIGS, "figS1_redside.png"), dpi=150)
    plt.close(fig)


def main():
    fig_redside()
    fig_residuals()
    fig_limits()
    fig_regions()
    fig_map()
    fig_environment()


if __name__ == "__main__":
    main()