"""Shared helpers: paths, downloads, great-circle distances, small statistics, plot style."""
import json
import os
import urllib.request

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
STUDY = os.path.abspath(os.path.join(HERE, ".."))
DATA = os.path.join(STUDY, "data")
CUR = os.path.join(DATA, "curated")
FIG = os.path.join(STUDY, "figures")
RAW = os.environ.get("EXPANSION_ANALOGS_RAW", "fermi_raw/expansion_analogs")
SEED = 20261012

# Raw public data files (re-downloaded on demand)
DOWNLOADS = {
    # GAVIA: Dyer, Redding & Blackburn 2017, Scientific Data 4, 170041; figshare 10.6084/m9.figshare.4234850
    "gavia": ("gavia/GAVIA_main_data_table.csv", "https://ndownloader.figshare.com/files/6940181"),
    # Moulton & Cropper 2019 PeerJ 7:e7637, supplementary data (Sol et al. 2012 introduction events), via Europe PMC
    "moulton_zip": ("propagule/pmc6744947_supp.zip",
                    "https://www.ebi.ac.uk/europepmc/webservices/rest/PMC6744947/supplementaryFiles"),
    # Cassey et al. 2018 PLoS Biol 16:e2005987, S1 Data (effect sizes of the propagule-pressure meta-analysis)
    "cassey": ("propagule/cassey2018_S1_Data.csv",
               "https://journals.plos.org/plosbiology/article/file?type=supplementary&id=10.1371/journal.pbio.2005987.s004"),
    # p3k14c global radiocarbon database (Bird et al. 2022, Scientific Data 9, 27)
    "p3k14c": ("p3k/p3k14c_scrubbed_fuzzed.csv",
               "https://raw.githubusercontent.com/people3k/p3k14c/main/inst/p3k14c_scrubbed_fuzzed.csv"),
}


def fetch(key, force=False):
    """Return the local path of a raw file, downloading it if it is missing."""
    rel, url = DOWNLOADS[key]
    path = os.path.join(RAW, rel)
    if force or not os.path.exists(path) or os.path.getsize(path) < 1000:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (research script; expansion_analogs)"})
        with urllib.request.urlopen(req, timeout=300) as r:
            blob = r.read()
        if len(blob) < 1000:
            raise IOError("download of %s returned only %d bytes" % (url, len(blob)))
        with open(path + ".part", "wb") as f:   # write to a temporary name so a failed download never truncates a good file
            f.write(blob)
        os.replace(path + ".part", path)
    return path


def haversine(lat1, lon1, lat2, lon2):
    """Great-circle distance in km (inputs in degrees; arrays broadcast)."""
    R = 6371.0088
    p1, p2 = np.radians(lat1), np.radians(lat2)
    dphi = p2 - p1
    dlmb = np.radians(np.asarray(lon2) - np.asarray(lon1))
    a = np.sin(dphi / 2) ** 2 + np.cos(p1) * np.cos(p2) * np.sin(dlmb / 2) ** 2
    return 2 * R * np.arcsin(np.sqrt(np.clip(a, 0, 1)))


def jeffreys(k, n, q=(0.025, 0.975)):
    """Jeffreys interval for a binomial proportion: posterior Beta(k+1/2, n-k+1/2)."""
    from scipy.stats import beta
    return tuple(float(beta.ppf(x, k + 0.5, n - k + 0.5)) for x in q)


def logistic_fit(X, y, ridge=0.0):
    """Maximum-likelihood logistic regression (optionally ridge-penalised on the slopes).
    X: (n, k) without intercept. Returns coefficients (intercept first), covariance, log-likelihood."""
    from scipy.optimize import minimize
    X = np.asarray(X, float)
    y = np.asarray(y, float)
    A = np.column_stack([np.ones(len(y)), X])

    def nll(b):
        z = A @ b
        return np.sum(np.logaddexp(0, z) - y * z) + 0.5 * ridge * np.sum(b[1:] ** 2)

    def grad(b):
        z = A @ b
        p = 1 / (1 + np.exp(-z))
        g = A.T @ (p - y)
        g[1:] += ridge * b[1:]
        return g

    res = minimize(nll, np.zeros(A.shape[1]), jac=grad, method="BFGS")
    b = res.x
    p = 1 / (1 + np.exp(-(A @ b)))
    H = (A * (p * (1 - p))[:, None]).T @ A + ridge * np.diag([0] + [1] * X.shape[1])
    try:
        cov = np.linalg.inv(H)
    except np.linalg.LinAlgError:
        cov = np.full_like(H, np.nan)
    ll = -np.sum(np.logaddexp(0, A @ b) - y * (A @ b))
    return b, cov, ll


def logistic_fit_fast(X, y, ridge=0.05, iters=25):
    """Ridge-penalised logistic regression by Newton steps (for bootstrap loops; same model as logistic_fit).
    Returns the coefficients, intercept first."""
    X = np.asarray(X, float)
    y = np.asarray(y, float)
    A = np.column_stack([np.ones(len(y)), X])
    P = ridge * np.diag([0.0] + [1.0] * X.shape[1])
    b = np.zeros(A.shape[1])
    for _ in range(iters):
        z = np.clip(A @ b, -30, 30)
        p = 1 / (1 + np.exp(-z))
        g = A.T @ (p - y) + P @ b
        H = (A * (p * (1 - p))[:, None]).T @ A + P + 1e-9 * np.eye(A.shape[1])
        step = np.linalg.solve(H, g)
        # damp very large steps (near-separation in a bootstrap resample)
        m = np.abs(step).max()
        if m > 5:
            step *= 5 / m
        b = b - step
        if m < 1e-6:
            break
    return b


def dump_json(obj, name):
    path = os.path.join(DATA, name)
    with open(path, "w") as f:
        json.dump(obj, f, indent=1, default=lambda o: float(o) if isinstance(o, (np.floating, np.integer)) else str(o))
    return path


# ---- plot style (light surface; categorical slots 1-3 validated for all-pairs use) ----
C = dict(surface="#fcfcfb", ink="#0b0b0b", ink2="#52514e", grid="#e4e3df", muted="#8a8983",
         blue="#2a78d6", orange="#eb6834", aqua="#1baf7a", yellow="#eda100", violet="#6250d6", gray="#a9a8a2")


def style():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "figure.facecolor": C["surface"], "axes.facecolor": C["surface"], "savefig.facecolor": C["surface"],
        "axes.edgecolor": C["grid"], "axes.labelcolor": C["ink2"], "text.color": C["ink"],
        "xtick.color": C["ink2"], "ytick.color": C["ink2"], "axes.grid": True, "grid.color": C["grid"],
        "grid.linewidth": 0.6, "axes.spines.top": False, "axes.spines.right": False,
        "font.size": 9.5, "axes.titlesize": 10.5, "axes.titleweight": "bold", "axes.titlelocation": "left",
        "legend.frameon": False, "legend.fontsize": 8.5, "lines.linewidth": 2.0, "figure.dpi": 150,
        "savefig.dpi": 170, "savefig.bbox": "tight", "font.family": "DejaVu Sans",
    })
    return plt