"""Band-averaged WISE fluxes of a blackbody enclosure (added in the revision).

A sphere that re-radiates the stellar luminosity L as a blackbody of temperature T has radius
R = sqrt(L / (4 pi sigma T^4)).  On the Wien tail (T < ~100 K at 22 um) the flux inside a WISE band
is dominated by the red edge of the passband, so a monochromatic evaluation at the isophotal
wavelength under-estimates it.  The catalogue converts counts to magnitudes with zero points
defined for F_nu ~ nu^-2 (Wright et al. 2010), so the flux density that the catalogue would
assign to a source of spectrum F_nu is

    F_cat = Int F_nu(lam) R(lam) dlam/lam  /  Int (lam/lam_iso)^2 R(lam) dlam/lam ,

with R(lam) the relative spectral response per photon (IRSA RSR-W?.txt).
"""
import os

import numpy as np

from common import RAW, H_PLANCK, C_LIGHT, K_B, SIGMA_SB, L_SUN, AU, PC

LAM_ISO = {"W1": 3.3526, "W2": 4.6028, "W3": 11.5608, "W4": 22.0883}      # um
F0_JY = {"W1": 309.540, "W2": 171.787, "W3": 31.674, "W4": 8.363}          # Jy, zero magnitude
_RSR = {}


def rsr(band):
    if band not in _RSR:
        a = np.loadtxt(os.path.join(RAW, f"RSR-{band}.txt"))
        _RSR[band] = (a[:, 0], np.clip(a[:, 1], 0, None))
    return _RSR[band]


def bnu(lam_um, T):
    """Planck function B_nu in W m^-2 Hz^-1 sr^-1 (overflow-safe)."""
    nu = C_LIGHT / (np.asarray(lam_um, float) * 1e-6)
    x = np.minimum(H_PLANCK * nu / (K_B * T), 700.0)
    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 band_bnu(band, T):
    """Catalogue-equivalent surface brightness: the Planck function averaged over the band as above."""
    lam, R = rsr(band)
    num = np.trapezoid(bnu(lam, T) * R / lam, lam)
    den = np.trapezoid((lam / LAM_ISO[band]) ** 2 * R / lam, lam)
    return num / den


def fcat_jy(band, T, L_lsun, d_pc, band_averaged=True):
    """Catalogue flux density [Jy] of a blackbody enclosure of luminosity L at distance d."""
    b = band_bnu(band, T) if band_averaged else bnu(LAM_ISO[band], T)
    R = radius_au(L_lsun, T) * AU
    return np.pi * b * (R / (d_pc * PC)) ** 2 / 1e-26


def dmax_pc(band, flim_jy, T, L_lsun, band_averaged=True):
    return 10.0 * np.sqrt(fcat_jy(band, T, L_lsun, 10.0, band_averaged) / flim_jy)


def colour(band_a, band_b, T, band_averaged=True):
    """Vega colour m_a - m_b of a blackbody at temperature T (large and positive when band a is dark)."""
    fa = band_bnu(band_a, T) if band_averaged else bnu(LAM_ISO[band_a], T)
    fb = band_bnu(band_b, T) if band_averaged else bnu(LAM_ISO[band_b], T)
    with np.errstate(divide="ignore"):
        return float(np.clip(-2.5 * np.log10(max(fa, 1e-300) / F0_JY[band_a]) + 2.5 * np.log10(fb / F0_JY[band_b]), -50, 99))