"""Part 0: how cold must a complete, isotropically radiating enclosure be to escape the
all-sky infrared surveys, how big is it then, and what does that size do to (i) binaries and
(ii) microlensing?

A sphere that re-radiates the full stellar luminosity L as a blackbody of temperature T has
radius R = sqrt(L / (4 pi sigma T^4)) and flux density F_nu = pi B_nu(T) (R/d)^2.

Survey limits (verified at source, see report):
  AllWISE SNR=5: 0.054, 0.071, 0.73, 5.0 mJy at 3.4, 4.6, 12, 22 um (AllWISE Expl. Suppl. II.3.a)
  IRAS PSC completeness: 0.4, 0.5, 0.6, 1.0 Jy at 12, 25, 60, 100 um (away from confusion)
  AKARI/FIS BSC: about 0.55 Jy at 90 um

Revision: (i) the WISE fluxes are also computed band-averaged over the published relative
spectral response (wiselib.py), because the threshold temperature sits on the Wien tail where a
monochromatic evaluation at 22.09 um under-estimates the in-band flux; the threshold is quoted as
the range between the two.  (ii) "Detectable" here means only "above the catalogue flux limit":
a 22-um-only source without an optical counterpart has no distance and is not thereby recognised
or excluded.  What the catalogue actually contains is counted in 02b_w4_dark_counts.py.
"""
import json
import os
import sys

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

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (DATA, FIGS, H_PLANCK, C_LIGHT, K_B, SIGMA_SB, L_SUN, AU, PC, KAPPA_MAS)
import wiselib

BANDS = {  # name: (wavelength um, limiting flux density Jy)
    "WISE W3 12um": (11.56, 0.73e-3),
    "WISE W4 22um": (22.09, 5.0e-3),
    "IRAS 25um": (25.0, 0.5),
    "IRAS 60um": (60.0, 0.6),
    "IRAS 100um": (100.0, 1.0),
    "AKARI 90um": (90.0, 0.55),
}


def bnu(lam_um, T):
    nu = C_LIGHT / (lam_um * 1e-6)
    x = H_PLANCK * nu / (K_B * T)
    return 2 * H_PLANCK * nu ** 3 / C_LIGHT ** 2 / np.expm1(x)


def radius_au(L_lsun, T):
    return np.sqrt(L_lsun * L_SUN / (4 * np.pi * SIGMA_SB * T ** 4)) / AU


def fnu_jy(lam_um, T, L_lsun, d_pc):
    R = radius_au(L_lsun, T) * AU
    return np.pi * bnu(lam_um, T) * (R / (d_pc * PC)) ** 2 / 1e-26


def dmax_pc(lam_um, flim_jy, T, L_lsun):
    return 10.0 * np.sqrt(fnu_jy(lam_um, T, L_lsun, 10.0) / flim_jy)


def einstein_radius_au(M, DL_kpc, DS_kpc=8.0):
    pirel = 1.0 / DL_kpc - 1.0 / DS_kpc           # mas
    return np.sqrt(KAPPA_MAS * M * pirel) * DL_kpc  # mas * kpc = au


def main():
    T = np.array([20, 25, 30, 40, 50, 75, 100, 150, 200, 300, 500, 1000], float)
    rows = []
    for L in (1.0, 0.1, 0.01):
        for t in T:
            r = {"L_Lsun": L, "T_K": t, "R_au": radius_au(L, t)}
            for name, (lam, flim) in BANDS.items():
                r["dmax_pc " + name] = dmax_pc(lam, flim, t, L)
            rows.append(r)
    tab = pd.DataFrame(rows)
    tab.to_csv(os.path.join(DATA, "ir_detectability.csv"), index=False, float_format="%.4g")

    # temperatures below which a sphere is undetectable beyond 100 pc and 20 pc
    Tfine = np.linspace(10, 400, 3901)
    summary = {}
    for L in (1.0, 0.1, 0.01):
        for dref in (20.0, 100.0, 300.0):
            for label, names in (("WISE", ["WISE W3 12um", "WISE W4 22um"]),
                                 ("WISE+IRAS+AKARI", list(BANDS))):
                d = np.max([dmax_pc(*BANDS[n], Tfine, L) for n in names], axis=0)
                ok = Tfine[d < dref]
                summary[f"Tmax_hidden_beyond_{dref:.0f}pc_L{L:g}_{label}"] = float(ok.max()) if len(ok) else None
    # band-averaged WISE fluxes (revision)
    wb = {"WISE W3 12um": "W3", "WISE W4 22um": "W4"}
    Tb = np.linspace(15, 120, 2101)
    for L in (1.0, 0.1, 0.01):
        for dref in (20.0, 100.0, 300.0):
            d = np.max([[wiselib.dmax_pc(wb[n], BANDS[n][1], t, L) for t in Tb] for n in wb], axis=0)
            ok = Tb[d < dref]
            summary[f"Tmax_hidden_beyond_{dref:.0f}pc_L{L:g}_WISE_bandavg"] = float(ok.max()) if len(ok) else None
    for L in (1.0, 0.01):
        for t in (100.0, 75.0, 50.0, 40.0, 30.0):
            summary[f"dmax_pc_W4_L{L:g}_T{t:.0f}_mono"] = float(dmax_pc(*BANDS["WISE W4 22um"], t, L))
            summary[f"dmax_pc_W4_L{L:g}_T{t:.0f}_bandavg"] = float(wiselib.dmax_pc("W4", BANDS["WISE W4 22um"][1], t, L))
    brows = []
    for L in (1.0, 0.1, 0.01):
        for t in T:
            brows.append({"L_Lsun": L, "T_K": t, "R_au": radius_au(L, t),
                          "dmax_pc W3 monochromatic": dmax_pc(*BANDS["WISE W3 12um"], t, L),
                          "dmax_pc W3 band-averaged": wiselib.dmax_pc("W3", BANDS["WISE W3 12um"][1], t, L),
                          "dmax_pc W4 monochromatic": dmax_pc(*BANDS["WISE W4 22um"], t, L),
                          "dmax_pc W4 band-averaged": wiselib.dmax_pc("W4", BANDS["WISE W4 22um"][1], t, L),
                          "W3-W4 band-averaged": wiselib.colour("W3", "W4", t),
                          "W2-W4 band-averaged": wiselib.colour("W2", "W4", t),
                          "W1-W4 band-averaged": wiselib.colour("W1", "W4", t)})
    pd.DataFrame(brows).to_csv(os.path.join(DATA, "ir_detectability_bandavg.csv"), index=False, float_format="%.4g")
    summary["R_au_L1_T300"] = float(radius_au(1, 300))
    summary["R_au_L1_T100"] = float(radius_au(1, 100))
    summary["R_au_L1_T50"] = float(radius_au(1, 50))
    summary["R_au_L1_T30"] = float(radius_au(1, 30))
    summary["Teq_at_1au_L1"] = float((L_SUN / (4 * np.pi * SIGMA_SB * AU ** 2)) ** 0.25)

    # microlensing: ratio of enclosure radius to Einstein radius; T at which they are equal
    ml = []
    for (M, L) in ((1.0, 1.0), (0.7, 0.15), (0.5, 0.04), (0.3, 0.011), (0.15, 0.003)):
        for DL in (2.0, 4.0, 6.0, 7.0):
            rE = einstein_radius_au(M, DL)
            Teq1 = summary["Teq_at_1au_L1"]
            Tcrit = Teq1 * (np.sqrt(L) / rE) ** 0.5
            ml.append({"M_Msun": M, "L_Lsun": L, "DL_kpc": DL, "rE_au": rE,
                       "rho_300K": radius_au(L, 300) / rE, "rho_100K": radius_au(L, 100) / rE,
                       "rho_50K": radius_au(L, 50) / rE, "rho_30K": radius_au(L, 30) / rE,
                       "T_crit_K (R_s = r_E)": Tcrit})
    mltab = pd.DataFrame(ml)
    mltab.to_csv(os.path.join(DATA, "microlens_occultation_geometry.csv"), index=False, float_format="%.4g")
    summary["Tcrit_K_M1_DL6"] = float(mltab[(mltab.M_Msun == 1.0) & (mltab.DL_kpc == 6.0)]["T_crit_K (R_s = r_E)"].iloc[0])
    summary["Tcrit_K_M0.3_DL6"] = float(mltab[(mltab.M_Msun == 0.3) & (mltab.DL_kpc == 6.0)]["T_crit_K (R_s = r_E)"].iloc[0])
    with open(os.path.join(DATA, "ir_detectability_summary.json"), "w") as f:
        json.dump(summary, f, indent=1)

    # figure
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.3))
    Tg = np.logspace(np.log10(15), np.log10(1000), 300)
    cols = {"WISE W3 12um": "#1f77b4", "WISE W4 22um": "#d62728", "IRAS 60um": "#2ca02c",
            "IRAS 100um": "#9467bd", "AKARI 90um": "#8c564b"}
    for name, c in cols.items():
        lam, flim = BANDS[name]
        if name in wb:      # WISE: band-averaged over the relative spectral response
            ax[0].loglog(Tg, [wiselib.dmax_pc(wb[name], flim, t, 1.0) for t in Tg], color=c, label=name + " (band-averaged)")
            ax[0].loglog(Tg, [wiselib.dmax_pc(wb[name], flim, t, 0.01) for t in Tg], color=c, ls=":")
            ax[0].loglog(Tg, dmax_pc(lam, flim, Tg, 1.0), color=c, lw=0.6, ls="--")
        else:
            ax[0].loglog(Tg, dmax_pc(lam, flim, Tg, 1.0), color=c, label=name + " (nominal catalogue limit)")
            ax[0].loglog(Tg, dmax_pc(lam, flim, Tg, 0.01), color=c, ls=":")
    ax[0].axhline(100, color="k", lw=0.6)
    ax[0].text(16, 115, "100 pc", fontsize=8)
    ax[0].set_ylim(1, 3e5)
    ax[0].set_xlabel("enclosure temperature T [K]")
    ax[0].set_ylabel("distance within which the source is\nabove the catalogue flux limit [pc]")
    ax[0].set_title("Blackbody enclosure re-radiating 1 (solid) or 0.01 (dotted) L$_\\odot$\n(thin dashed: WISE evaluated monochromatically)", fontsize=9)
    ax[0].legend(fontsize=8, loc="lower right")
    ax[0].grid(alpha=0.3, which="both")

    ax[1].loglog(Tg, radius_au(1.0, Tg), "k", label="enclosure radius, L = 1 L$_\\odot$")
    ax[1].loglog(Tg, radius_au(0.01, Tg), "k:", label="enclosure radius, L = 0.01 L$_\\odot$")
    ax[1].axhspan(einstein_radius_au(0.3, 6.0), einstein_radius_au(1.0, 4.0), color="#ff7f0e", alpha=0.3,
                  label="Einstein radius (0.3-1 M$_\\odot$ lens at 4-6 kpc)")
    ax[1].axhspan(0.2, 1.0, color="#1f77b4", alpha=0.25,
                  label="largest stable enclosure in a P = 100-1000 d binary")
    ax[1].set_xlabel("enclosure temperature T [K]")
    ax[1].set_ylabel("radius [au]")
    ax[1].set_ylim(0.01, 1e3)
    ax[1].legend(fontsize=8)
    ax[1].grid(alpha=0.3, which="both")
    ax[1].set_title("Size of a complete enclosure versus the two gravitational scales", fontsize=10)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig1_ir_detectability.png"), dpi=150)
    plt.close(fig)
    print("[ir] ", json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()