"""Part 1a: the signal.  Size of an opaque blackbody enclosure, when it hides from WISE, how the
size is distributed over the stellar population, and the characteristic occultation time scales.

Outputs: data/radius_table.csv, data/wise_hidden_temperature.csv, data/population_radii.csv,
         data/signal_summary.json, figures/fig1_signal.png
"""
import json
import os

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

from common import (DATA, FIGS, H_PLANCK, C_LIGHT, K_B, AU, PC, R_SUN, T1AU, KMS_TO_AU_PER_YR, radius_au)
import popmodel as pm

# AllWISE SNR=5 point-source limits (Explanatory Supplement II.3.a), as used in studies/dark_companions
WISE = {"W3 12um": (11.56, 0.73e-3), "W4 22um": (22.09, 5.0e-3)}   # wavelength um, Jy


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


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


def t_hidden(L, d_pc):
    """Temperature below which the enclosure is fainter than both AllWISE limits at distance d
    (nominal, uncrowded-sky limits).  Returns inf if it is never brighter than the limits below 600 K."""
    T = np.linspace(8, 600, 5921)
    det = np.zeros_like(T, bool)
    for lam, lim in WISE.values():
        det |= fnu_jy(lam, T, L, d_pc) > lim
    return float(T[np.argmax(det)]) if det.any() else np.inf


def main():
    Ls = [1e-3, 3e-3, 1e-2, 3e-2, 0.1, 0.3, 1.0]
    Ts = [10, 20, 30, 50, 75, 100, 150, 200, 300]
    rows = []
    for L in Ls:
        for T in Ts:
            R = float(radius_au(L, T))
            rows.append({"L_Lsun": L, "T_K": T, "R_au": R, "diameter_au": 2 * R,
                         # mean chord (pi/2 R) crossing time at two reference transverse speeds
                         "mean_duration_d_at_200kms": 0.5 * np.pi * R / (200 * KMS_TO_AU_PER_YR) * 365.25,
                         "mean_duration_d_at_40kms": 0.5 * np.pi * R / (40 * KMS_TO_AU_PER_YR) * 365.25})
    tab = pd.DataFrame(rows)
    tab.to_csv(os.path.join(DATA, "radius_table.csv"), index=False, float_format="%.4g")
    piv = tab.pivot(index="T_K", columns="L_Lsun", values="R_au")
    print("[signal] enclosure radius R [au]; rows T [K], columns L [Lsun]")
    print(piv.to_string(float_format=lambda x: f"{x:.3g}"))

    hid = []
    for L in Ls:
        for d in (100.0, 1000.0, 6000.0):
            th = t_hidden(L, d)
            hid.append({"L_Lsun": L, "d_pc": d, "T_hidden_K": th,
                        "R_au_at_T_hidden": float(radius_au(L, th)) if np.isfinite(th) else np.nan})
    hid = pd.DataFrame(hid)
    hid.to_csv(os.path.join(DATA, "wise_hidden_temperature.csv"), index=False, float_format="%.4g")
    print(hid.pivot(index="L_Lsun", columns="d_pc", values="T_hidden_K").to_string(float_format=lambda x: f"{x:.0f}"))

    # population: radii for a Kroupa PDMF (0.08-1 Msun) with the Mamajek mass-luminosity relation
    norm = pm.ms_stars_per_msun()
    prow = []
    for T in (20, 30, 50, 75, 100, 150, 200, 300):
        m, w, R = pm.population_radii(T)
        o = np.argsort(R)
        cdf = np.cumsum(w[o])
        q = lambda p: float(np.interp(p, cdf, R[o]))
        prow.append({"T_K": T, "R_mean_au": float(np.sum(w * R)), "R_median_au": q(0.5), "R_10pct_au": q(0.1),
                     "R_90pct_au": q(0.9), "R2_mean_au2": float(np.sum(w * R ** 2)),
                     "frac_stars_R_gt_3au": float(w[R > 3].sum()), "R_sun_au": float(radius_au(1.0, T))})
    pop = pd.DataFrame(prow)
    pop.to_csv(os.path.join(DATA, "population_radii.csv"), index=False, float_format="%.4g")
    print(pop.to_string(index=False, float_format=lambda x: f"{x:.3g}"))

    # ingress time: source radius crossing at relative speed v (lens-plane projection ~ source radius for
    # a lens well in front of a bulge source; exact factor D_l/D_s <= 1)
    ingress = {f"{name} ({rs} Rsun) at {v} km/s [h]": 2 * rs * R_SUN / (v * 1e3) / 3600.0
               for name, rs in (("dwarf", 1.0), ("red-clump giant", 10.0), ("bright giant", 50.0)) for v in (40, 200)}
    summary = {
        "T_1au_K": T1AU,
        "R_au_1Lsun_30K": float(radius_au(1.0, 30.0)),
        "diameter_au_1Lsun_30K": 2 * float(radius_au(1.0, 30.0)),
        "brief_claim_checked": "A 1 Lsun enclosure at 30 K has R = 172 au (344 au across), not 9 au; "
                               "9 au across corresponds to 1 Lsun at 184 K or to 7e-4 Lsun at 30 K.",
        "T_for_9au_diameter_1Lsun_K": float(T1AU / np.sqrt(4.5)),
        "L_for_9au_diameter_30K_Lsun": float((4.5 / (T1AU / 30.0) ** 2) ** 2),
        "population_normalisation": norm,
        "mean_sqrtL": float(np.sum(pm.mass_grid()[1] * np.sqrt(pm.lum_of_mass(pm.mass_grid()[0])))),
        "ingress_times_hours": ingress,
        "T_hidden_from_WISE_K": {f"L={r.L_Lsun:g} d={r.d_pc:g}pc": (r.T_hidden_K if np.isfinite(r.T_hidden_K) else
                                                                    "never detectable below 600 K") for r in hid.itertuples()},
    }
    with open(os.path.join(DATA, "signal_summary.json"), "w") as f:
        json.dump(summary, f, indent=1)
    print(json.dumps(summary, indent=1))

    # ---- figure 1 --------------------------------------------------------------------------------
    fig, ax = plt.subplots(1, 2, figsize=(12.5, 4.8))
    Tg = np.logspace(1, np.log10(300), 200)
    cols = plt.cm.viridis(np.linspace(0.05, 0.9, len(Ls)))
    for L, c in zip(Ls, cols):
        ax[0].loglog(Tg, radius_au(L, Tg), color=c, lw=1.6, label=f"{L:g}")
        th = t_hidden(L, 100.0)
        if th < 300:
            ax[0].plot([th], [radius_au(L, th)], "o", color=c, ms=5, mec="k", mew=0.6)
        th6 = t_hidden(L, 6000.0)
        if th6 < 300:
            ax[0].plot([th6], [radius_au(L, th6)], "s", color=c, ms=5, mec="k", mew=0.6)
    ax[0].axhline(2.0, color="0.5", ls=":", lw=1)
    ax[0].text(10.5, 2.25, "typical Einstein radius of a bulge lens (below: lensing, no occultation)", fontsize=7.5, color="0.3")
    ax[0].set_xlabel("enclosure temperature T [K]")
    ax[0].set_ylabel("enclosure radius R [au]")
    ax[0].legend(title="L [L$_\\odot$]", fontsize=8, title_fontsize=8, ncol=2, loc="upper right")
    ax[0].set_title("Size of an opaque blackbody enclosure\n(circles/squares: coldest T detectable by AllWISE at 100 pc / 6 kpc)", fontsize=9)
    ax[0].set_xlim(10, 300)
    ax[0].set_ylim(0.02, 3000)
    ax[0].grid(alpha=0.3, which="both")
    sec = ax[0].secondary_yaxis("right", functions=(lambda r: 0.5 * np.pi * r / (200 * KMS_TO_AU_PER_YR) * 365.25,
                                                    lambda d: d / 365.25 * (200 * KMS_TO_AU_PER_YR) / (0.5 * np.pi)))
    sec.set_ylabel("mean occultation duration at 200 km/s [d]")

    for T, c in zip((30, 50, 100, 200), ("#1f77b4", "#2ca02c", "#ff7f0e", "#d62728")):
        m, w, R = pm.population_radii(T)
        o = np.argsort(R)
        ax[1].semilogx(R[o], np.cumsum(w[o]), color=c, lw=1.6, label=f"T = {T} K")
    ax[1].axvline(2.0, color="0.5", ls=":", lw=1)
    ax[1].set_xlabel("enclosure radius R [au]")
    ax[1].set_ylabel("cumulative fraction of main-sequence stars (0.08-1 M$_\\odot$)")
    ax[1].set_title("Most stars are M dwarfs: their enclosures are 10-30 times\nsmaller than a solar one at the same temperature", fontsize=9)
    ax[1].legend(fontsize=8)
    ax[1].grid(alpha=0.3, which="both")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "fig1_signal.png"), dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()