"""SFD (Schlegel, Finkbeiner & Davis 1998) E(B-V) lookup from the LAMBDA HEALPix map
(nside 512, NESTED, Galactic coordinates), with a small pure-numpy ang2pix."""
import os

import numpy as np
from astropy.io import fits

from common import RAW, http_get

SFD_URL = "https://lambda.gsfc.nasa.gov/data/foregrounds/SFD/lambda_sfd_ebv.fits"
SFD_FILE = os.path.join(RAW, "lambda_sfd_ebv.fits")
_MAP = None


def _spread(v):
    v = v.astype(np.int64)
    out = np.zeros_like(v)
    for i in range(16):
        out |= ((v >> i) & 1) << (2 * i)
    return out


def ang2pix_nest(nside, lon_deg, lat_deg):
    order = int(np.log2(nside))
    z = np.sin(np.radians(np.asarray(lat_deg, dtype=float)))
    za = np.abs(z)
    tt = (np.asarray(lon_deg, dtype=float) % 360.0) / 90.0
    # equatorial
    t1 = nside * (0.5 + tt)
    t2 = nside * z * 0.75
    jp = np.floor(t1 - t2).astype(np.int64)
    jm = np.floor(t1 + t2).astype(np.int64)
    ifp, ifm = jp >> order, jm >> order
    face_e = np.where(ifp == ifm, ifp | 4, np.where(ifp < ifm, ifp, ifm + 8))
    ix_e = jm & (nside - 1)
    iy_e = nside - (jp & (nside - 1)) - 1
    # polar
    ntt = np.minimum(3, tt.astype(np.int64))
    tp = tt - ntt
    tmp = nside * np.sqrt(3 * (1 - za))
    jpp = np.minimum((tp * tmp).astype(np.int64), nside - 1)
    jmp = np.minimum(((1.0 - tp) * tmp).astype(np.int64), nside - 1)
    north = z >= 0
    ix_p = np.where(north, nside - jmp - 1, jpp)
    iy_p = np.where(north, nside - jpp - 1, jmp)
    face_p = np.where(north, ntt, ntt + 8)
    eq = za <= 2.0 / 3.0
    ix = np.where(eq, ix_e, ix_p)
    iy = np.where(eq, iy_e, iy_p)
    face = np.where(eq, face_e, face_p)
    return face * nside * nside + _spread(ix) + (_spread(iy) << 1)


def sfd_ebv(glon, glat):
    global _MAP
    if _MAP is None:
        if not os.path.exists(SFD_FILE) or os.path.getsize(SFD_FILE) < 2e7:
            with open(SFD_FILE, "wb") as f:
                f.write(http_get(SFD_URL, post=False))
        with fits.open(SFD_FILE) as h:
            _MAP = np.array(h[1].data["TEMPERATURE"], dtype=float)
    return _MAP[ang2pix_nest(512, glon, glat)]


if __name__ == "__main__":
    tests = {"NGP": (0, 90), "SGP": (0, -90), "Gal centre": (0, 0), "LMC": (280.47, -32.89),
             "Orion Neb": (209.01, -19.38), "Taurus TMC1": (174.1, -13.7), "rho Oph": (353.7, 17.7),
             "Lockman hole": (149.8, 52.0), "anticentre b=30": (180, 30), "l=90 b=0": (90, 0)}
    for k, (l, b) in tests.items():
        print(f"{k:18s} l={l:7.2f} b={b:6.2f} E(B-V)={sfd_ebv(np.array([l]), np.array([b]))[0]:.3f}")
    # pixel-count sanity: uniform random points must fill pixels uniformly
    rng = np.random.default_rng(1)
    l = rng.uniform(0, 360, 2000000); b = np.degrees(np.arcsin(rng.uniform(-1, 1, 2000000)))
    p = ang2pix_nest(8, l, b); c = np.bincount(p, minlength=768)
    print("nside 8: pixels hit", (c > 0).sum(), "of 768; count min/max", c.min(), c.max())