"""Shared helpers for the tf_kiii study (paths, TAP download with caching)."""
import os
import time
import urllib.parse
import urllib.request

STUDY = "studies/tf_kiii"
RAW = "fermi_raw/tf_kiii"
DATA = os.path.join(STUDY, "data")
FIGS = os.path.join(STUDY, "figures")
for _d in (RAW, DATA, FIGS, os.path.join(RAW, "cutouts"), os.path.join(RAW, "cones")):
    os.makedirs(_d, exist_ok=True)

VIZIER_TAP = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"
IRSA_TAP = "https://irsa.ipac.caltech.edu/TAP/sync"
H0 = 70.0          # km/s/Mpc, ALFALFA convention (Haynes et al. 2018)
C_KMS = 299792.458


def http_get(url, data=None, timeout=300, retries=4):
    last = None
    for k in range(retries):
        try:
            req = urllib.request.Request(url, data=data,
                                         headers={"User-Agent": "tf_kiii-study/1.0"})
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as e:  # noqa
            last = e
            time.sleep(3 * (k + 1))
    raise RuntimeError(f"download failed: {url[:120]} : {last}")


def tap_csv(endpoint, query, outfile, fmt="csv", force=False, maxrec=None):
    """Run a synchronous TAP query, cache the CSV in RAW, return the path."""
    path = os.path.join(RAW, outfile)
    if os.path.exists(path) and os.path.getsize(path) > 0 and not force:
        return path
    params = {"REQUEST": "doQuery", "LANG": "ADQL", "FORMAT": fmt, "QUERY": query}
    if maxrec:
        params["MAXREC"] = str(maxrec)
    blob = http_get(endpoint, data=urllib.parse.urlencode(params).encode())
    if blob[:200].lstrip().startswith(b"<"):
        raise RuntimeError(f"TAP error for {outfile}: {blob[:600]!r}")
    with open(path, "wb") as f:
        f.write(blob)
    return path


# ---------------------------------------------------------------- utilities
import numpy as np  # noqa: E402


def radec_to_xyz(ra_deg, dec_deg):
    ra = np.radians(np.asarray(ra_deg, float))
    dec = np.radians(np.asarray(dec_deg, float))
    return np.c_[np.cos(dec) * np.cos(ra), np.cos(dec) * np.sin(ra), np.sin(dec)]


def chord_for_arcsec(arcsec):
    return 2 * np.sin(np.radians(arcsec / 3600.0) / 2)


def mad_sigma(x):
    x = np.asarray(x, float)
    return 1.4826 * np.median(np.abs(x - np.median(x)))


def robust_polyfit(x, y, deg=2, clip=2.5, niter=20):
    """Iteratively clipped least-squares polynomial; returns (coeffs, robust sigma)."""
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    keep = np.isfinite(x) & np.isfinite(y)
    sig = np.nan
    p = None
    for _ in range(niter):
        p = np.polyfit(x[keep], y[keep], deg)
        r = y - np.polyval(p, x)
        sig = mad_sigma(r[keep])
        new = np.isfinite(r) & (np.abs(r) < clip * sig)
        if new.sum() == keep.sum() and np.all(new == keep):
            break
        keep = new
    return p, sig


def poisson_ul(n, cl=0.95):
    """Classical one-sided Poisson upper limit on the mean given n observed."""
    from scipy.stats import chi2
    return 0.5 * chi2.ppf(cl, 2 * (n + 1))


def cmb_velocity(ra_deg, dec_deg, vhel):
    """Heliocentric -> CMB frame (Planck 2018 dipole: 369.82 km/s to l=264.021, b=48.253)."""
    from astropy.coordinates import SkyCoord
    import astropy.units as u
    g = SkyCoord(np.asarray(ra_deg) * u.deg, np.asarray(dec_deg) * u.deg).galactic
    l, b = g.l.rad, g.b.rad
    l0, b0, v0 = np.radians(264.021), np.radians(48.253), 369.82
    return np.asarray(vhel) + v0 * (np.sin(b) * np.sin(b0) + np.cos(b) * np.cos(b0) * np.cos(l - l0))


def running_median_fit(x, y, nbins=15):
    """Non-parametric relation: median of y in nbins equal-count bins of x.

    Returns a function f(x) that interpolates linearly between the bin medians and extrapolates
    linearly beyond the first and last bin centres (using the slope of the outermost two bins).
    Used as an alternative to the clipped quadratic (fit-form sensitivity check)."""
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    ok = np.isfinite(x) & np.isfinite(y)
    xs, ys = x[ok], y[ok]
    edges = np.quantile(xs, np.linspace(0, 1, nbins + 1))
    idx = np.clip(np.searchsorted(edges, xs, side="right") - 1, 0, nbins - 1)
    xc = np.array([np.median(xs[idx == k]) for k in range(nbins)])
    yc = np.array([np.median(ys[idx == k]) for k in range(nbins)])

    def f(xq):
        xq = np.asarray(xq, float)
        out = np.interp(xq, xc, yc)
        lo, hi = xq < xc[0], xq > xc[-1]
        out = np.where(lo, yc[0] + (xq - xc[0]) * (yc[1] - yc[0]) / (xc[1] - xc[0]), out)
        out = np.where(hi, yc[-1] + (xq - xc[-1]) * (yc[-1] - yc[-2]) / (xc[-1] - xc[-2]), out)
        return out
    return f


def bb_frac(lam_um, T):
    """nu L_nu / L_bol of a blackbody of temperature T (K) at wavelength lam (micron)."""
    x = 14387.8 / (np.asarray(lam_um, float) * np.asarray(T, float))
    return (15 / np.pi ** 4) * x ** 4 / np.expm1(x)


# ---- constants shared by the waste-heat calculations
LAM_W1 = 3.368            # micron, WISE W1 isophotal wavelength
NULNU_SUN_I = 0.70        # nu L_nu of the Sun at SDSS i, in Lsun (M_i,AB = 4.53)
NULNU_SUN_W1 = 0.0425     # nu L_nu of the Sun at W1, in Lsun (M_W1,AB = 5.94 = 3.24 Vega + 2.699)
BC_I = 2.0                # adopted L_bol(stars) / nu L_nu(i) of a disc galaxy (an assumption; see the write-up, waste-heat temperature)


def clopper_pearson(k, n, cl=0.68):
    """Two-sided binomial interval (lo, hi) for k successes out of n."""
    from scipy.stats import beta
    a = (1 - cl) / 2
    lo = 0.0 if k == 0 else beta.ppf(a, k, n - k + 1)
    hi = 1.0 if k == n else beta.ppf(1 - a, k + 1, n - k)
    return lo, hi