"""Partial Dyson-swarm photometric model.

A swarm with covering fraction f intercepts f of the stellar luminosity (grey, so every
stellar band dims by 1-f) and re-radiates f*L as a blackbody of temperature T.  The
swarm's flux density at Earth is  F_nu = f * F_bol,* * pi B_nu(T) / (sigma T^4).

WISE magnitudes are band-integrated, so the *apparent* flux density (F0 * 10^(-m/2.5))
of a blackbody differs from its monochromatic value at the isophotal wavelength by the
flux-correction factor f_c (Wright et al. 2010, Table 1).  We compute f_c from the WISE
relative system response curves (photon-counting: signal ~ int F_lambda * lambda * R dlambda),
normalised to the Vega-calibrated convention (f_c = 1.0084, 1.0066, 1.0088, 1.0013 for
F_nu ~ nu^2).  This reproduces the published f_c for power laws to <0.1%.
"""
import os

import numpy as np

from common import RAW, LAM, SIGMA_SB, planck_nu, http_get

RSR_URL = "https://irsa.ipac.caltech.edu/data/WISE/docs/release/Prelim/expsup/figures/RSR-{b}.txt"
FC_RJ = dict(W1=1.0084, W2=1.0066, W3=1.0088, W4=1.0013)
_RSR = {}


def rsr(b):
    if b not in _RSR:
        path = os.path.join(RAW, f"RSR-{b}.txt")
        if not os.path.exists(path) or os.path.getsize(path) < 2000:
            with open(path, "wb") as f:
                f.write(http_get(RSR_URL.format(b=b), post=False))
        a = np.loadtxt(path)
        _RSR[b] = (a[:, 0], a[:, 1])
    return _RSR[b]


def flux_correction(b, fnu):
    """f_c for a spectrum fnu(lambda_um) in band b (apparent = f_c * monochromatic)."""
    lam, R = rsr(b)

    def sig(f):
        return np.trapezoid(f(lam) / lam ** 2 * lam * R, lam)

    def ref(x):
        return x ** -2.0

    return FC_RJ[b] * (sig(fnu) / fnu(LAM[b])) / (sig(ref) / ref(LAM[b]))


def swarm_response(T):
    """Apparent flux density (Jy) per unit swarm bolometric flux (W m^-2) in each band."""
    out = {}
    for b in ("J", "H", "Ks", "W1", "W2", "W3", "W4"):
        mono = np.pi * planck_nu(LAM[b], T) / (SIGMA_SB * T ** 4) * 1e26   # Jy per (W m^-2)
        if b.startswith("W"):
            mono = mono * flux_correction(b, lambda x: planck_nu(x, T))
        out[b] = float(mono)
    return out


if __name__ == "__main__":
    print("check vs Wright+2010 / Explanatory Supplement IV.4.h: F_nu=const -> f_c(W3)=0.9169;"
          " F_nu~nu^-1 -> 0.992, 0.994, 0.937, 0.993")
    print("  const:", round(flux_correction("W3", lambda x: x ** 0.0), 4))
    print("  nu^-1:", [round(flux_correction(b, lambda x: x ** 1.0), 4) for b in ("W1", "W2", "W3", "W4")])
    for T in (100, 150, 200, 300, 400, 600):
        r = swarm_response(T)
        print(T, {k: f"{v:.3e}" for k, v in r.items()},
              "fc:", [round(flux_correction(b, lambda x: planck_nu(x, T)), 3) for b in ("W1", "W2", "W3", "W4")])