"""Shared helpers for the mir_quiescent study (mid-infrared waste-heat limits for quiescent galaxies)."""
import io
import os
import time
import urllib.parse
import urllib.request
import uuid

import numpy as np
from scipy import stats

HERE = os.path.dirname(os.path.abspath(__file__))
STUDY = os.path.dirname(HERE)
RAW = os.environ.get("MIRQ_RAW", "fermi_raw/mir_quiescent")
DATA = os.path.join(STUDY, "data")
FIGS = os.path.join(STUDY, "figures")
CACHE = os.path.join(RAW, "derived")      # large re-creatable tables stay out of the shared folder
for _d in (RAW, DATA, FIGS, CACHE):
    os.makedirs(_d, exist_ok=True)

C_KMS = 299792.458
H0 = 70.0
OM = 0.3

# model grid required by the study brief
F_GRID = np.array([0.01, 0.03, 0.1, 0.3])
T_GRID = np.array([100.0, 150.0, 200.0, 300.0, 400.0, 600.0])
# finer grid used for figures / the limit curve
F_FINE = np.array([0.003, 0.005, 0.01, 0.02, 0.03, 0.05, 0.07, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4, 0.5])
# grid on which every analysis variant is run (brief grid plus the f values used in comparisons with prior work)
F_VAR = np.array([0.01, 0.03, 0.1, 0.2, 0.25, 0.3, 0.5])
# temperature bands for envelopes: the project model's warm band and the band where WISE is most sensitive
T_BAND_FULL = (100.0, 150.0, 200.0, 300.0, 400.0, 600.0)
T_BAND_CORE = (150.0, 200.0, 300.0, 400.0)

# WISE zero-magnitude flux densities (Jy; Jarrett et al. 2011) and coadd instrumental zero points (mag)
WISE_F0 = {1: 309.540, 2: 171.787, 3: 31.674, 4: 8.363}
WISE_ZP = {1: 20.5, 2: 19.5, 3: 18.0, 4: 13.0}
WISE_LAM_ISO = {1: 3.3526, 2: 4.6028, 3: 11.5608, 4: 22.0883}      # micron
NMGY_JY = 3.631e-6                                                  # 1 nanomaggy (AB) in Jy
# Vega -> AB offsets used by the Legacy Surveys for their WISE forced photometry
LS_VEGA_TO_AB = {1: 2.699, 2: 3.339, 3: 5.174, 4: 6.620}


def wise_dn_to_jy(dn, band):
    """AllWISE profile-fit 'raw' flux (w?flux, DN) -> Jy with the catalogue zero points."""
    return WISE_F0[band] * 10 ** (-0.4 * WISE_ZP[band]) * np.asarray(dn, float)


def wise_mag_to_jy(mag, band):
    return WISE_F0[band] * 10 ** (-0.4 * np.asarray(mag, float))


def ls_nmgy_to_wise_jy(nmgy, band):
    """Legacy Surveys WISE flux (AB nanomaggies) -> Jy on the WISE Vega-system scale used above."""
    vega_mag = 22.5 - 2.5 * np.log10(1.0) - LS_VEGA_TO_AB[band]         # mag of 1 nmgy in Vega
    return WISE_F0[band] * 10 ** (-0.4 * vega_mag) * np.asarray(nmgy, float)


# ----------------------------------------------------------------------------
# cosmology (flat LCDM)
# ----------------------------------------------------------------------------
_zg = np.linspace(0.0, 0.3, 3001)
_Ez = np.sqrt(OM * (1 + _zg) ** 3 + 1 - OM)
_dc = np.concatenate([[0.0], np.cumsum(0.5 * (1 / _Ez[1:] + 1 / _Ez[:-1]) * np.diff(_zg))]) * C_KMS / H0


def luminosity_distance(z):
    """Mpc"""
    return np.interp(z, _zg, _dc) * (1.0 + np.asarray(z))


# ----------------------------------------------------------------------------
# statistics
# ----------------------------------------------------------------------------
def poisson_ul(n, cl=0.95):
    return stats.chi2.ppf(cl, 2 * (np.asarray(n) + 1)) / 2.0


def binom_ul(k, n, cl=0.95):
    """One-sided Clopper-Pearson upper limit on a binomial proportion (k successes in n trials)."""
    k = np.asarray(k, float)
    n = np.asarray(n, float)
    with np.errstate(invalid="ignore"):
        out = np.where(k >= n, 1.0, stats.beta.ppf(cl, k + 1, np.maximum(n - k, 1e-9)))
    return np.where(n > 0, out, np.nan)


# ----------------------------------------------------------------------------
# network helpers
# ----------------------------------------------------------------------------
def fetch(url, dest, data=None, headers=None, timeout=900, tries=6, min_bytes=50, check=None):
    """Download url -> dest (cached).  `check(blob)` may raise to reject a payload."""
    if os.path.exists(dest) and os.path.getsize(dest) >= min_bytes:
        return dest
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    last = None
    for k in range(tries):
        try:
            req = urllib.request.Request(url, data=data, headers=headers or {})
            with urllib.request.urlopen(req, timeout=timeout) as r:
                blob = r.read()
            if len(blob) < min_bytes:
                raise RuntimeError("short payload: %r" % blob[:200])
            if check is not None:
                check(blob)
            with open(dest + ".part", "wb") as f:
                f.write(blob)
            os.replace(dest + ".part", dest)
            return dest
        except Exception as e:  # noqa
            last = e
            time.sleep(10 * (k + 1))
    raise RuntimeError("failed to fetch %s: %s" % (url[:200], last))


def multipart(fields, files):
    boundary = uuid.uuid4().hex
    body = io.BytesIO()
    for k, v in fields.items():
        body.write(("--%s\r\nContent-Disposition: form-data; name=\"%s\"\r\n\r\n%s\r\n" % (boundary, k, v)).encode())
    for k, (fn, blob) in files.items():
        body.write(("--%s\r\nContent-Disposition: form-data; name=\"%s\"; filename=\"%s\"\r\n"
                    "Content-Type: text/csv\r\n\r\n" % (boundary, k, fn)).encode())
        body.write(blob)
        body.write(b"\r\n")
    body.write(("--%s--\r\n" % boundary).encode())
    return body.getvalue(), {"Content-Type": "multipart/form-data; boundary=%s" % boundary}


def not_html_or_error(blob):
    head = blob[:400].lower()
    if b"<html" in head or b"error" in head[:120] or head[:1] == b"{":
        raise RuntimeError("bad payload: %r" % blob[:200])


def angsep_arcsec(ra1, dec1, ra2, dec2):
    r1, d1, r2, d2 = map(np.radians, (ra1, dec1, ra2, dec2))
    s = np.sin((d2 - d1) / 2) ** 2 + np.cos(d1) * np.cos(d2) * np.sin((r2 - r1) / 2) ** 2
    return np.degrees(2 * np.arcsin(np.sqrt(np.clip(s, 0, 1)))) * 3600.0