"""Empirical photosphere model: predicted W3 / W4 fluxes from Ks and from W1, with an
empirical error model.  Used identically on the real stars and on injected swarms.

Calibration (done once on the real sample, stored in a dict `cal`):
  cK1, cK3, cK4 : median Ks-W1, Ks-W3, Ks-W4 colour of clean, high-S/N stars, as a function
                  of TWO predictors: the Gaia absolute magnitude M_G (suffix _M) and the Gaia
                  colour G-RP (suffix _C).  The prediction uses the mean of the two and half
                  their difference is added to the error budget.  (M_G fails for pre-main-
                  sequence stars and for hosts dimmed by a swarm; G-RP fails for close pairs
                  with contaminated RP fluxes; where the two agree the extra term vanishes.)
  d1(W1)        : correction of the AllWISE profile-fit W1 magnitude of saturated stars
                  (W1 < 8.3) onto the Ks system (cf. Patel et al. 2014), and the empirical
                  scatter e1sat(W1) of the corrected magnitudes.
  d3(W3pred)    : the same for bright W3.
  s3K, s3W, s4K, s4W : fractional error floor (intrinsic colour scatter + calibration)
                  versus predicted magnitude for each band/anchor, chosen so that the robust
                  width of the significance distribution is 1.
"""
import numpy as np

from common import mag2jy, jy2mag

XM_BINS = np.concatenate([np.arange(0.5, 6.0, 0.75), np.arange(6.0, 12.0, 0.4), np.arange(12.0, 17.01, 0.75)])
XC_BINS = np.concatenate([np.arange(-0.1, 0.5, 0.1), np.arange(0.5, 1.5, 0.04), np.arange(1.5, 1.81, 0.06)])
W1_SAT = 8.3


def robust_sigma(x):
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x)]
    if len(x) < 5:
        return np.nan
    return 1.4826 * np.median(np.abs(x - np.median(x)))


def binned_locus(x, y, bins, nmin=12):
    """Median of y in bins of x -> (x_nodes, y_nodes) for np.interp (flat extrapolation)."""
    xs, ys = [], []
    ok = np.isfinite(x) & np.isfinite(y)
    x, y = x[ok], y[ok]
    idx = np.digitize(x, bins)
    for k in range(1, len(bins)):
        sel = idx == k
        if sel.sum() >= nmin:
            xs.append(np.median(x[sel]))
            ys.append(np.median(y[sel]))
    return np.array(xs), np.array(ys)


def interp(nodes, x):
    return np.interp(x, nodes[0], nodes[1])


def clean_mask(d):
    """Stars usable for calibrating the photospheric colour loci."""
    return ((d.cc_flags.str[0] == "0") & (d.cc_flags.str[2] == "0") & (d.cc_flags.str[3] == "0")
            & (d.nb == 1) & (d.na == 0) & (d.n_aw_neigh == 0) & (d.ext_flg <= 1) & (d.n_share == 1)
            & (d.ebv < 0.15) & (d.e_Ksmag < 0.03) & (d.Ksmag > 4.0)).values


def colours(cal, M_G, g_rp):
    """Mean of the two predictors and half their difference, for Ks-W1, Ks-W3, Ks-W4."""
    out = {}
    for k in ("cK1", "cK3", "cK4"):
        a = interp(cal[k + "_M"], M_G)
        b = interp(cal[k + "_C"], g_rp)
        out[k] = 0.5 * (a + b)
        out["h" + k] = 0.5 * (a - b)
    return out


def calibrate(d):
    """d: DataFrame of matched stars with columns M_G, g_rp, Ksmag, e_Ksmag, FW1.., flags."""
    cal = {}
    xm, xc = d.M_G.values, d.g_rp.values
    ks = d.Ksmag.values
    with np.errstate(invalid="ignore", divide="ignore"):
        w1 = jy2mag(d.FW1.values, "W1")
        w3 = jy2mag(d.FW3.values, "W3")
        w4 = jy2mag(d.FW4.values, "W4")
    cl = clean_mask(d)
    # loci are fitted on stars within 1 mag of the main-sequence ridge so that the two
    # predictors describe the same population
    ms = cl & (np.abs(d.dM.values) < 1.0)
    # --- Ks-W1 locus from unsaturated W1
    s = ms & (w1 > W1_SAT) & (w1 < 12.5) & np.isfinite(w1)
    cal["cK1_M"] = binned_locus(xm[s], (ks - w1)[s], XM_BINS)
    cal["cK1_C"] = binned_locus(xc[s], (ks - w1)[s], XC_BINS)
    cK1 = 0.5 * (interp(cal["cK1_M"], xm) + interp(cal["cK1_C"], xc))
    # --- saturated-W1 correction onto the Ks system, and its empirical scatter
    r1 = ks - w1 - cK1
    s = ms & (w1 <= 8.6) & np.isfinite(w1)
    b1 = np.arange(2.0, 8.61, 0.3)
    nodes = binned_locus(w1[s], r1[s], b1, nmin=10)
    nodes = (np.append(nodes[0][nodes[0] < W1_SAT], [W1_SAT, 30.0]), np.append(nodes[1][nodes[0] < W1_SAT], [0.0, 0.0]))
    cal["d1"] = nodes
    w1c = w1 + interp(cal["d1"], w1)
    r1c = ks - w1c - cK1
    xs, ys = [], []
    for a, b in zip(b1[:-1], b1[1:]):
        sel = s & (w1 >= a) & (w1 < b)
        if sel.sum() >= 10:
            v = robust_sigma(r1c[sel]) ** 2 - np.median(d.e_Ksmag.values[sel]) ** 2
            xs.append(0.5 * (a + b))
            ys.append(0.921 * np.sqrt(max(v, 0.02 ** 2)))
    cal["e1sat"] = (np.array(xs), np.array(ys))
    # --- Ks-W3 locus (unsaturated W3, high S/N)
    s = ms & (d.w3snr.values > 15) & (w3 > 4.5) & np.isfinite(w3)
    cal["cK3_M"] = binned_locus(xm[s], (ks - w3)[s], XM_BINS)
    cal["cK3_C"] = binned_locus(xc[s], (ks - w3)[s], XC_BINS)
    cK3 = 0.5 * (interp(cal["cK3_M"], xm) + interp(cal["cK3_C"], xc))
    m3p = ks - cK3
    r3 = ks - w3 - cK3
    s = ms & (m3p <= 5.1) & np.isfinite(w3)
    nodes = binned_locus(m3p[s], r3[s], np.arange(1.5, 5.11, 0.4), nmin=8)
    if len(nodes[0]):
        nodes = (np.append(nodes[0][nodes[0] < 4.7], [4.7, 30.0]), np.append(nodes[1][nodes[0] < 4.7], [0.0, 0.0]))
    else:
        nodes = (np.array([0.0, 30.0]), np.array([0.0, 0.0]))
    cal["d3"] = nodes
    # --- Ks-W4 locus (W4 S/N > 10: bright nearby stars only); beyond the calibrated range
    #     W3-W4 is held constant (the W4 photosphere of such stars is far below the noise)
    s = ms & (d.w4snr.values > 10) & np.isfinite(w4)
    for suf, x, bins in (("_M", xm, XM_BINS), ("_C", xc, XC_BINS)):
        k4 = binned_locus(x[s], (ks - w4)[s], bins, nmin=10)
        xk3 = cal["cK3" + suf][0]
        red = xk3 > k4[0].max()
        c34_last = k4[1][-1] - np.interp(k4[0][-1], *cal["cK3" + suf])
        cal["cK4" + suf] = (np.append(k4[0], xk3[red]), np.append(k4[1], cal["cK3" + suf][1][red] + c34_last))
    # median 2MASS Ks uncertainty versus Ks (used to assign an uncertainty to a dimmed host)
    cal["eKs_curve"] = binned_locus(ks, d.e_Ksmag.values, np.arange(4.0, 16.01, 0.5), nmin=20)
    return cal


def predict(cal, M_G, g_rp, FKs, FW1):
    """Predicted photospheric W3, W4 fluxes (Jy) from each anchor, plus the fractional
    uncertainty of each prediction that comes from the disagreement of the two predictors."""
    c = colours(cal, np.asarray(M_G, dtype=float), np.asarray(g_rp, dtype=float))
    with np.errstate(invalid="ignore", divide="ignore"):
        ks = jy2mag(FKs, "Ks")
        w1 = jy2mag(FW1, "W1")
    w1 = np.where(np.isfinite(w1), w1, 30.0)
    w1c = w1 + interp(cal["d1"], w1)
    out = {}
    m3K = ks - c["cK3"]
    m3K = m3K - interp(cal["d3"], m3K)          # predicted *observed* W3 magnitude
    m3W = w1c - (c["cK3"] - c["cK1"])
    m3W = m3W - interp(cal["d3"], m3W)
    m4K = ks - c["cK4"]
    m4W = w1c - (c["cK4"] - c["cK1"])
    out["P3K"], out["P3W"] = mag2jy(m3K, "W3"), mag2jy(m3W, "W3")
    out["P4K"], out["P4W"] = mag2jy(m4K, "W4"), mag2jy(m4W, "W4")
    out["h3K"], out["h3W"] = 0.921 * np.abs(c["hcK3"]), 0.921 * np.abs(c["hcK3"] - c["hcK1"])
    out["h4K"], out["h4W"] = 0.921 * np.abs(c["hcK4"]), 0.921 * np.abs(c["hcK4"] - c["hcK1"])
    return out


def anchor_errors(cal, e_Ks, FW1, eW1):
    """Fractional errors of the two anchors (Ks; saturation-corrected W1)."""
    eK = 0.921 * np.asarray(e_Ks, dtype=float)
    FW1 = np.asarray(FW1, dtype=float)
    e1 = np.asarray(eW1, dtype=float) / np.clip(FW1, 1e-9, None)
    with np.errstate(invalid="ignore", divide="ignore"):
        w1 = jy2mag(FW1, "W1")
    sat = w1 < W1_SAT
    if len(cal["e1sat"][0]):
        e1 = np.where(sat, np.interp(w1, *cal["e1sat"]), e1)
    return eK, e1


def _floor_nodes(mpred, frac_resid, var_other_frac, good, lo, hi, step=0.5, nmin=40, smin=0.01, max_other=0.06):
    """Fractional error floor vs predicted magnitude: s^2 = robust_var(frac resid) - median(other var).
    Only bins where the other (measurement + anchor) errors are small enough to see the floor."""
    xs, ys = [], []
    edges = np.arange(lo, hi + 1e-9, step)
    for a, b in zip(edges[:-1], edges[1:]):
        sel = good & (mpred >= a) & (mpred < b)
        if sel.sum() < nmin:
            continue
        other = np.median(var_other_frac[sel])
        if not np.isfinite(other) or other > max_other ** 2:
            continue
        v = robust_sigma(frac_resid[sel]) ** 2 - other
        xs.append(0.5 * (a + b))
        ys.append(np.sqrt(max(v, smin ** 2)))
    if not xs:
        xs, ys = [8.0], [0.03]
    return np.array(xs), np.array(ys)


def fit_error_floor(cal, d, pred):
    """Fit fractional floors s3K, s3W, s4K, s4W (functions of predicted magnitude)."""
    good = ((d.cc_flags.str[0] == "0") & (d.nb == 1) & (d.na == 0) & (d.n_aw_neigh == 0)
            & (d.ext_flg <= 1) & (d.n_share == 1) & (d.ebv < 0.15)).values
    eK, e1 = anchor_errors(cal, d.e_Ksmag.values, d.FW1.values, d.eW1.values)
    for band, F, eF, flagpos in (("3", d.FW3.values, d.eW3.values, 2), ("4", d.FW4.values, d.eW4.values, 3)):
        g = good & (d.cc_flags.str[flagpos] == "0").values
        for anc, ea in (("K", eK), ("W", e1)):
            P = pred[f"P{band}{anc}"]
            mp = jy2mag(P, "W" + band)
            frac = (F - P) / P
            other = (eF / P) ** 2 + ea ** 2 + pred[f"h{band}{anc}"] ** 2
            cal[f"s{band}{anc}"] = _floor_nodes(mp, frac, other, g & np.isfinite(frac), 1.0, 12.0)
    return cal


def significance(cal, e_Ks, FW1, eW1, F3, e3, F4, e4, pred):
    """Excess flux and significance for both anchors; returns dict of arrays."""
    eK, e1 = anchor_errors(cal, e_Ks, FW1, eW1)
    out = {}
    for band, F, eF in (("3", F3, e3), ("4", F4, e4)):
        for anc, ea in (("K", eK), ("W", e1)):
            P = pred[f"P{band}{anc}"]
            mp = jy2mag(P, "W" + band)
            s = np.interp(mp, *cal[f"s{band}{anc}"])
            sig = np.sqrt(eF ** 2 + P ** 2 * (ea ** 2 + s ** 2 + pred[f"h{band}{anc}"] ** 2))
            out[f"E{band}{anc}"] = F - P
            out[f"sig{band}{anc}"] = sig
            out[f"chi{band}{anc}"] = (F - P) / sig
        out[f"chi{band}"] = np.minimum(out[f"chi{band}K"], out[f"chi{band}W"])
        out[f"chi{band}_neg"] = np.maximum(out[f"chi{band}K"], out[f"chi{band}W"])  # mirror statistic
    return out