#!/usr/bin/env python3
"""Figures (PNG) for the mir_quiescent study.  Reads only the tables in data/ (and the master table)."""
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import mirq_common as M
import s05_limits as L
from s03_model import Model

# palette (validated categorical slots 1-3 + neutrals); text always in ink colours
BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
INK, INK2, GRID, SURF, GREY = "#0b0b0b", "#52514e", "#e6e5e1", "#fcfcfb", "#9a9992"
T_RAMP = ["#86b6ef", "#5598e7", "#2a78d6", "#1c5cab", "#104281", "#0d366b"]       # ordinal one-hue ramp

plt.rcParams.update({
    "figure.facecolor": SURF, "axes.facecolor": SURF, "savefig.facecolor": SURF,
    "axes.edgecolor": GREY, "axes.labelcolor": INK, "text.color": INK, "xtick.color": INK2, "ytick.color": INK2,
    "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6, "axes.spines.top": False,
    "axes.spines.right": False, "font.size": 9, "axes.titlesize": 9.5, "legend.frameon": False,
    "lines.linewidth": 1.8, "axes.axisbelow": True,
})


def fig1_baseline(d, fl, m):
    """Natural baseline: distribution of W3/W1 and W4/W1 for bright galaxies, with model predictions."""
    P = L.phot(d, "psf")
    clean = L.clean_mask(fl, True)
    br = P["ok"] & (P["F1"] > 2.0)
    zmed = float(np.median(d.zhelio.values[br]))
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.0))
    for ax, b, T in ((axes[0], 3, 300.0), (axes[1], 4, 150.0)):
        r = P["F%d" % b][br] / P["F1"][br]
        rc = P["F%d" % b][br & clean] / P["F1"][br & clean]
        lo, hi = (0.05, 30.0) if b == 3 else (0.02, 100.0)
        bins = np.logspace(np.log10(lo), np.log10(hi), 70)
        ax.hist(np.clip(r, lo * 1.01, hi * 0.99), bins=bins, histtype="stepfilled", color=BLUE, alpha=0.25, lw=0)
        ax.hist(np.clip(r, lo * 1.01, hi * 0.99), bins=bins, histtype="step", color=BLUE, lw=1.4,
                label="all early types (N=%d)" % br.sum())
        ax.hist(np.clip(rc, lo * 1.01, hi * 0.99), bins=bins, histtype="step", color=ORANGE, lw=1.4,
                label="clean quiescent (N=%d)" % (br & clean).sum())
        c = float(m.C_of(b, zmed))
        ax.axvline(c, color=INK2, lw=1.0)
        ax.text(c, 1.02, "old stars\n(Brown+14)", transform=ax.get_xaxis_transform(), ha="center", va="bottom",
                fontsize=7, color=INK2)
        R1, Rb = m.R_of(T, 1, zmed), m.R_of(T, b, zmed)
        for f in M.F_GRID:
            g = f / (1 - f)
            x = (c + g * Rb) / (1 + g * R1)
            ax.axvline(x, color=AQUA, lw=1.0)
            ax.text(x, 0.97, "f=%g" % f, transform=ax.get_xaxis_transform(), rotation=90, ha="right", va="top",
                    fontsize=7.5, color=INK2)
        ax.set_xscale("log")
        ax.set_yscale("log")
        ax.set_xlim(lo, hi)
        ax.set_xlabel("W%d / W1 flux-density ratio (AllWISE profile fit)" % b)
        ax.set_ylabel("galaxies per bin")
        ax.set_title("W%d: a host of a %d K radiator would sit at the green lines (z=%.2f)" % (b, T, zmed), pad=22)
    h, lab = axes[0].get_legend_handles_labels()
    fig.legend(h, lab, loc="upper center", ncol=2, fontsize=8, bbox_to_anchor=(0.5, 1.0))
    fig.suptitle("Natural mid-infrared baseline of bright (W1 > 2 mJy) early-type galaxies "
                 "(first and last bins hold the overflow)", y=1.06, fontsize=10.5)
    fig.tight_layout()
    fig.savefig(os.path.join(M.FIGS, "fig1_baseline.png"), dpi=160, bbox_inches="tight")
    plt.close(fig)


def fig2_sensitivity():
    t = pd.read_csv(os.path.join(M.DATA, "limits_primary.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 3.9))
    for T, col in zip(M.T_GRID, T_RAMP):
        s = t[t["T"] == T].sort_values("f")
        axes[0].plot(s.f, s.sens_frac, color=col, marker="o", ms=3.5, label="%d K" % T)
        axes[1].plot(s.f, np.clip(s.frac_compatible, 1e-6, 1), color=col, marker="o", ms=3.5, label="%d K" % T)
    axes[0].set_xscale("log")
    axes[0].set_xlabel("covering fraction f")
    axes[0].set_ylabel("fraction of galaxies in which the signal is testable")
    axes[0].set_title("Sensitivity fraction (catalogue-level injection; model-based)")
    axes[0].legend(title="waste-heat T", fontsize=7.5, title_fontsize=7.5, loc="upper left")
    axes[1].set_xscale("log")
    axes[1].set_yscale("log")
    axes[1].set_xlabel("covering fraction f")
    axes[1].set_ylabel("candidates / testable galaxies")
    axes[1].set_title("Natural confusion: testable galaxies that look like hosts")
    fig.tight_layout()
    fig.savefig(os.path.join(M.FIGS, "fig2_sensitivity.png"), dpi=160, bbox_inches="tight")
    plt.close(fig)


def fig3_limits():
    """Own limits only: one panel per temperature."""
    t = pd.read_csv(os.path.join(M.DATA, "limits_primary.csv"))
    vfn = os.path.join(M.DATA, "limits_vetted.csv")
    v = pd.read_csv(vfn) if os.path.exists(vfn) else None
    var = pd.read_csv(os.path.join(M.DATA, "limits_variants.csv"))
    _E4 = var[var.label == "eps0.4"]
    fig, axes = plt.subplots(2, 3, figsize=(11.5, 7.0), sharex=True, sharey=True)
    for ax, T in zip(axes.ravel(), M.T_GRID):
        s = t[t["T"] == T].sort_values("f")
        ax.plot(s.f, s.ul_zero, color=GREY, lw=1.2, label="sensitivity floor (if there were no candidates)")
        ax.plot(s.f, s.ul_all, color=BLUE, marker="o", ms=3.5, label="all outliers, all early types")
        e4 = _E4[_E4["T"] == T].sort_values("f")
        ax.plot(e4.f, e4.ul_all, color=BLUE, lw=1.0, ls=(0, (4, 2)), label="all outliers, conservative tolerance 0.4")
        ax.plot(s.f, s.ul_clean, color=ORANGE, marker="s", ms=3.5, label="clean quiescent subsample (catalogue-unexplained)")
        if v is not None:
            sv = v[(v["T"] == T) & (v.n_cand_not_inspected == 0)].sort_values("f")      # fully inspected cells only
            ax.plot(sv.f, sv.ul_unexplained, color=AQUA, marker="^", ms=6, lw=0, mec=SURF, mew=0.6,
                    label="clean, after image vetting (fully inspected cells)")
        ax.set_xscale("log")
        ax.set_yscale("log")
        ax.set_xlim(0.0025, 0.6)
        ax.set_ylim(1e-5, 1.5)
        ax.set_title("T = %d K" % T)
    for ax in axes[1]:
        ax.set_xlabel("covering fraction f")
    for ax in axes[:, 0]:
        ax.set_ylabel("95% upper limit on fraction of galaxies")
    h, lab = axes[0, 0].get_legend_handles_labels()
    fig.legend(h, lab, loc="upper center", ncol=2, fontsize=8, bbox_to_anchor=(0.5, 1.07))
    fig.tight_layout()
    fig.savefig(os.path.join(M.FIGS, "fig3_limits.png"), dpi=160, bbox_inches="tight")
    plt.close(fig)


def fig4_comparison():
    """This work (300 K; weakest limit over the 100-600 K band; conservative tolerance) against published limits."""
    t = pd.read_csv(os.path.join(M.DATA, "limits_primary.csv"))
    h = pd.read_csv(os.path.join(M.DATA, "headline_limits.csv"))
    fig, ax = plt.subplots(figsize=(8.0, 5.6))
    s3 = t[t["T"] == 300.0].sort_values("f")
    ax.plot(s3.f, s3.ul_all, color=BLUE, marker="o", ms=3.5, label="this work, 300 K, all outliers")
    ax.plot(h.f, h["ul_all_full100_600K"], color=BLUE, lw=1.2, ls=(0, (4, 2)),
            label="this work, weakest of 100-600 K, all outliers")
    ax.fill_between(h.f, h["ul_all_full100_600K"], h["ul_all_full100_600K_eps0.4"], color=BLUE, alpha=0.15, lw=0,
                    zorder=0, label="the same with model tolerance 0.2 to 0.4")
    ax.plot(s3.f, s3.ul_clean, color=ORANGE, marker="s", ms=3.5, label="this work, 300 K, clean quiescent subsample")
    ax.plot(h.f, h["ul_clean_full100_600K"], color=ORANGE, lw=1.2, ls=(0, (4, 2)),
            label="this work, weakest of 100-600 K, clean quiescent")
    # published points (see report for what each one means)
    ax.plot([0.05, 0.10, 0.25], [0.146, 0.101, 0.026], color=INK2, marker="x", ms=6, lw=0.8)
    ax.annotate("Curtis+26 (129 galaxies, 300 K)", (0.25, 0.026), xytext=(8, 6), textcoords="offset points",
                fontsize=7.5, color=INK2)
    ax.plot([0.5, 0.85], [3.0 / 4.8e3, 3.0 / 1e5], color=INK2, marker="v", ms=5.5, lw=0.8)
    ax.annotate("Griffith+15 as used in the\nproject model (gamma > 0.5, 0.85)", (0.85, 3.0 / 1e5), xytext=(-4, -22),
                textcoords="offset points", fontsize=7.5, color=INK2, ha="center")
    ax.plot([0.21, 0.6], [1.61e-4, 1.61e-4], color=INK2, marker="D", ms=4, lw=0.8, markevery=[0])
    ax.annotate("Huang+26 (2MRS, fixed L*, 300 K,\nAGN vetoed)", (0.6, 1.61e-4), xytext=(5, 0), textcoords="offset points",
                fontsize=7.5, color=INK2, ha="left", va="center")
    ax.plot([0.01, 0.1, 1.0], [1.25e-3, 3.3e-5, 1.2e-6], color=GREY, lw=1.2)
    ax.annotate("Lacki 16: cold (near-CMB) shrouds,\nPlanck; not a warm-waste-heat limit", (0.012, 1.0e-3),
                xytext=(0, -8), textcoords="offset points", fontsize=7.5, color=INK2, va="top", rotation=-31)
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlim(0.008, 3.0)
    ax.set_ylim(1e-6, 1.5)
    ax.set_xticks([0.01, 0.03, 0.1, 0.3, 1.0])
    ax.set_xticklabels(["0.01", "0.03", "0.1", "0.3", "1"])
    ax.set_xlabel("covering fraction f (share of starlight reprocessed)")
    ax.set_ylabel("95% upper limit on fraction of galaxies")
    ax.set_title("Warm waste heat: this work (early-type galaxies only) and published limits")
    ax.legend(loc="lower left", fontsize=7.2)
    fig.tight_layout()
    fig.savefig(os.path.join(M.FIGS, "fig4_comparison.png"), dpi=160, bbox_inches="tight")
    plt.close(fig)


def fig5_profile_fit_response():
    """Image-level test: does profile fitting return the modelled band ratios for star-following emission?"""
    fn = os.path.join(M.DATA, "image_injection_galaxies.csv")
    if not os.path.exists(fn):
        print("fig5 skipped: no image-level test table")
        return
    import json
    t = pd.read_csv(fn)
    g = pd.read_csv(os.path.join(M.DATA, "psf_response_grid.csv"))
    summ = json.load(open(os.path.join(M.DATA, "image_injection_summary.json")))
    fig, axes = plt.subplots(1, 3, figsize=(13.0, 4.0))
    # (a) emulator against the catalogue
    ax = axes[0]
    for b, col, snr, mk in ((1, BLUE, 20, "o"), (3, ORANGE, 10, "s"), (4, AQUA, 8, "^")):
        ok = t["cat%d" % b] / t["cat_e%d" % b] > snr
        ax.plot(t.theta_e[ok], (t["emu%d" % b] / t["cat%d" % b])[ok], mk, ms=3.2, color=col, alpha=0.65, mew=0,
                label="W%d (catalogue S/N > %d, n = %d)" % (b, snr, ok.sum()))
    ax.axhline(1.0, color=INK2, lw=0.8)
    ax.set_xscale("log")
    ax.set_ylim(0.5, 1.6)
    ax.set_xlabel("SDSS effective radius (arcsec)")
    ax.set_ylabel("emulated profile-fit flux / AllWISE catalogue flux")
    ax.set_title("(a) Emulated fit on unWISE images vs AllWISE")
    ax.legend(fontsize=7.2, loc="upper left")
    # (b) response to injected star-following sources
    ax = axes[1]
    for b, col, mk in ((3, ORANGE, "s"), (4, AQUA, "^")):
        ax.plot(t.theta_e, t["ratio_anch%d" % b], mk, ms=3.2, color=col, alpha=0.65, mew=0,
                label="W%d / W1, anchored to the galaxy's AllWISE W1" % b)
        gg = g[g.ba == 0.7]
        ax.plot(gg.devrad * np.sqrt(0.7), gg["eta%d" % b] / gg.eta1, color=col, lw=1.4,
                label="W%d / W1, emulator only (b/a = 0.7)" % b)
    ax.axhline(1.0, color=INK2, lw=0.8)
    ax.axhline(0.8, color=INK2, lw=0.8, ls=(0, (4, 2)))
    ax.text(0.55, 0.785, "primary tolerance: a host must reach 0.8", fontsize=7.2, color=INK2, va="top")
    ax.set_xscale("log")
    ax.set_xlim(0.5, 30)
    ax.set_ylim(0.6, 1.9)
    ax.set_xlabel("SDSS effective radius (arcsec)")
    ax.set_ylabel("measured band ratio / modelled band ratio")
    ax.set_title("(b) Injected waste heat that follows the starlight")
    ax.legend(fontsize=7.0, loc="upper left")
    # (c) W1 loss: whole sample against the response model
    ax = axes[2]
    c = pd.DataFrame(summ["w1_loss_check"])
    x = np.sqrt(c.theta_lo * np.minimum(c.theta_hi, 12.0))
    ax.plot(x, c.observed_F1pf_over_L1total, color=BLUE, marker="o", ms=4.5, label="observed: AllWISE profile-fit W1 / Legacy total W1")
    ax.plot(x, c.predicted_eta1, color=INK2, marker="D", ms=4, lw=1.0, label="predicted profile-fit response (de Vaucouleurs)")
    ax.set_xscale("log")
    ax.set_ylim(0, 1.05)
    ax.set_xlabel("SDSS effective radius (arcsec), bin centre")
    ax.set_ylabel("fraction of total W1 flux returned by profile fitting")
    ax.set_title("(c) W1 flux lost by profile fitting, whole sample")
    ax.legend(fontsize=7.2, loc="lower left")
    fig.tight_layout()
    fig.savefig(os.path.join(M.FIGS, "fig5_profile_fit_response.png"), dpi=160, bbox_inches="tight")
    plt.close(fig)


def main():
    d = L.load_master()
    fl = L.flags(d)
    m = Model()
    fig1_baseline(d, fl, m)
    fig2_sensitivity()
    fig3_limits()
    fig4_comparison()
    fig5_profile_fit_response()
    print("figures written to", M.FIGS)


if __name__ == "__main__":
    main()