"""Light-curve metrics and selection cuts for the search for total occultations in the OGLE-II
bulge DIA catalogue of candidate variables (Wozniak et al. 2002; 221,801 light curves, 49 fields,
difference-image flux at a fixed centroid, 1997-2000).

A totally occulted star shows a flat-bottomed drop of its *total* flux (reference + difference) to
the level of any unresolved blend, for as long as the occulter covers it.  Because the photometry
is forced difference-image flux, a star that vanishes is still measured (negative difference flux),
unlike in detection-based catalogues.

The difference flux is measured relative to a reference image that is a stack of about 20
good-seeing frames; `reference_weights` recovers which frames those were from the light curves
themselves (the stack-weighted mean difference flux of every object is zero by construction).
"""
import os
import warnings

import numpy as np
from astropy.io import fits

from common import RAW

DIA = os.path.join(RAW, "ogle2_dia")
MAPS = os.path.join(RAW, "ogle2_maps")
T_END_DETECTION = 1560.0      # HJD-2450000: end of the 1999 season; variables were detected on 1997-1999 frames
LOW_FRAC = 0.5                # "low" = total flux below this fraction of the bright level
CUTS = dict(i_max=18.0, idrad_max=1.5, nrun_min=5, dur_min=3.0, nlow_out_frac=0.02, lowfrac_max=0.6,
            edge_gap=30.0, edge_level=0.75, quiet_floor=0.05, quiet_nsig=2.5, flat_floor=0.07, flat_nsig=2.5)


ERR_CODE = -99.0              # pipeline error code for the difference flux and its error
QUALITY_BITS = 0x7FF          # light-curve flags 1-11 (catalog_description.txt, Appendix B): unreliable measurement.
                              # Flags 12-13 only say that no magnitude exists (total flux <= 0) and are kept.


def quality_mask(F, E, fl):
    """The one and only definition of a usable photometric point, applied to real and to injected
    light curves alike.  It must not depend on the value of the flux, apart from the error code:
    a star that vanishes has difference flux = -(its reference flux), i.e. hundreds to thousands
    of counts below zero.  (The first version of this study used `F > -90` here, which deleted
    exactly those points and made the search blind; see the report, section "Loader defect".)"""
    return ((fl & QUALITY_BITS) == 0) & (F != ERR_CODE) & (E > 0)


def load_field_raw(i):
    """Return catalogue recarray, time-sorted raw arrays (t, F, E, flags) and the frame table
    (names, JD) in the same order, for field i."""
    cat = fits.open(os.path.join(DIA, f"bul_sc{i}_cat4.fts"))[1].data
    with fits.open(os.path.join(DIA, f"bul_sc{i}_db4.fts")) as h:
        frames = np.array(h[1].data["FRAME"]).astype(str)
        jd = np.array(h[1].data["JD_HEL"], float)
        nobs = len(jd)
        db = h[2].data
        n = len(cat)
        t = np.array(db["OBS_TIME"], float).reshape(n, nobs)
        F = np.array(db["DIF_FLUX"], float).reshape(n, nobs)
        E = np.array(db["FLUX_ERR"], float).reshape(n, nobs)
        fl = np.array(db["FLAG"]).astype(np.int64).reshape(n, nobs)
    o = np.argsort(t[0])
    return cat, t[:, o], F[:, o], E[:, o], fl[:, o], np.char.strip(frames[o]), jd[o]


def load_field(i):
    """Return catalogue recarray and time-sorted arrays (t, F, E, good) for field i."""
    cat, t, F, E, fl, _, _ = load_field_raw(i)
    return cat, t, F, E, quality_mask(F, E, fl)


def getvar_like(t, F, E, good, sn_factor=0.7):
    """Emulation of the pixel-variability test of the DIA variable finder (Wozniak 2000, sec. 3.9),
    applied to the catalogue light curve on the 1997-1999 frames: the baseline is the median flux;
    the object is variable if at least 3 consecutive points depart from it by >= 3 sigma in the same
    direction, or at least 10 points (not necessarily consecutive) by >= 4 sigma in the same
    direction.  The real test runs on 3x3-pixel aperture light curves, whose signal-to-noise is
    about 0.86 of that of PSF photometry for 1.3 arcsec seeing; sn_factor = 0.7 is used to be safe.
    (Catalogue errors already contain the factor 1.6 that the variable finder used.)"""
    g = good & (t < T_END_DETECTION)
    if g.sum() < 20:
        return False
    z = (F[g] - np.median(F[g])) / E[g] * sn_factor
    for sgn in (1.0, -1.0):
        d3 = (sgn * z >= 3.0).astype(np.int8)
        if d3.size >= 3 and np.any(d3[2:] + d3[1:-1] + d3[:-2] == 3):
            return True
        if (sgn * z >= 4.0).sum() >= 10:
            return True
    return False


def metrics(t, F, E, good, ref, zp, low_frac=LOW_FRAC):
    """Metrics of the deepest sustained low state of one light curve (arrays over epochs)."""
    g = good
    ng = int(g.sum())
    if ng < 30:
        return None
    tg, Fg, Eg = t[g], F[g], E[g]
    B = np.percentile(Fg, 85)
    tot = ref + B
    if not tot > 0:
        return None
    r = (ref + Fg) / tot
    er = Eg / tot
    low = r < low_frac
    nlow = int(low.sum())
    if nlow < 3:
        return None
    d = np.diff(np.concatenate([[0], low.astype(np.int8), [0]]))
    st, en = np.where(d == 1)[0], np.where(d == -1)[0]
    j = int(np.argmax(en - st))
    a, b = int(st[j]), int(en[j])
    nrun = b - a
    rr = r[a:b]
    out = np.ones(ng, bool)
    out[a:b] = False
    ro = r[out]
    pre = r[max(0, a - 3):a]
    post = r[b:b + 3]
    return dict(
        ngood=ng, B=B, tot=tot, I_bright=zp - 2.5 * np.log10(tot), nlow=nlow, nrun=nrun,
        t1=tg[a], t2=tg[b - 1], dur=tg[b - 1] - tg[a],
        tprev=tg[a - 1] if a > 0 else np.nan, tnext=tg[b] if b < ng else np.nan,
        n_run_det=int((tg[a:b] < T_END_DETECTION).sum()),
        r_run_med=float(np.median(rr)), r_run_std=float(rr.std()), er_run=float(np.median(er[a:b])),
        pre_med=float(np.median(pre)) if pre.size else np.nan, post_med=float(np.median(post)) if post.size else np.nan,
        n_out=int(out.sum()), r_out_med=float(np.median(ro)) if ro.size else np.nan,
        mad_out=float(1.4826 * np.median(np.abs(ro - np.median(ro)))) if ro.size > 3 else np.nan,
        er_out=float(np.median(er[out])) if ro.size else np.nan,
        low_frac_pts=nlow / ng)


def cut_flags(m, idrad, c=CUTS):
    """Ordered dictionary of boolean cuts for one metrics dict."""
    ingress = (np.isfinite(m["tprev"]) and (m["t1"] - m["tprev"]) <= c["edge_gap"] and m["pre_med"] > c["edge_level"])
    egress = (np.isfinite(m["tnext"]) and (m["tnext"] - m["t2"]) <= c["edge_gap"] and m["post_med"] > c["edge_level"])
    return {
        "C1 bright state I<18": m["I_bright"] < c["i_max"],
        "C2 centroid on a reference-image star (ID_RAD<=1.5 px)": idrad <= c["idrad_max"],
        "C3 run of >=5 consecutive low points lasting >=3 d, >=3 of them in 1997-1999":
            (m["nrun"] >= c["nrun_min"]) and (m["dur"] >= c["dur_min"]) and (m["n_run_det"] >= 3),
        "C4 single episode (<=2% of points low outside the run) and <=60% of points low":
            ((m["nlow"] - m["nrun"]) <= max(1, c["nlow_out_frac"] * m["ngood"])) and (m["low_frac_pts"] <= c["lowfrac_max"]),
        "C5 ingress or egress resolved within a season (gap<=30 d, bright side >75%)": bool(ingress or egress),
        "C6 quiet outside the run": m["mad_out"] <= max(c["quiet_floor"], c["quiet_nsig"] * m["er_out"]),
        "C7 flat bottom": m["r_run_std"] <= max(c["flat_floor"], c["flat_nsig"] * m["er_run"]),
    }


CUT_NAMES = list(cut_flags(dict(I_bright=0, nrun=0, dur=0, n_run_det=0, nlow=0, ngood=1, low_frac_pts=0, tprev=np.nan,
                                tnext=np.nan, t1=0, t2=0, pre_med=0, post_med=0, mad_out=0, er_out=0, r_run_std=0,
                                er_run=0), 0).keys())


def passes_all(m, idrad):
    return all(cut_flags(m, idrad).values())


def reference_weights(t, F, E, good, jd, idr):
    """Recover the weights of the individual frames in the reference image of one field.

    By construction the difference flux of every object, averaged over the frames stacked into the
    reference image, is zero:  sum_j w_j F_ij = 0 for all objects i.  With hundreds of strongly
    variable objects per field this is a heavily over-determined non-negative least-squares problem
    for the weights w_j (>= 0, summing to 1) of the 1997-1999 frames (the reference images were not
    rebuilt when the 2000 season was added).  Each light curve is scaled by its own standard
    deviation, so that the residual is the stack-weighted mean flux in units of the variability
    amplitude; objects whose residual is more than 4 robust sigma out are dropped and the fit is
    repeated.  Returns the weights over all epochs (zero for 2000), the number of objects used, the
    rms residual and, for comparison, the rms of the unweighted mean over all 1997-1999 frames."""
    from scipy.optimize import nnls
    det = jd < T_END_DETECTION
    Fm = np.where(good, F, np.nan)
    with np.errstate(all="ignore"), warnings.catch_warnings():
        warnings.simplefilter("ignore")
        medE = np.nanmedian(np.where(good, E, np.nan), axis=1)
        sd = np.nanstd(Fm[:, det], axis=1)
    comp = good[:, det].mean(axis=1)
    rows = np.where((sd > 5 * medE) & (comp > 0.9) & np.isfinite(medE) & (idr <= CUTS["idrad_max"]))[0]
    A = F[rows][:, det] / sd[rows][:, None]
    g = good[rows][:, det]
    med = np.nanmedian(np.where(g, A, np.nan), axis=1)
    A = np.where(g, A, med[:, None])                           # impute the few flagged points
    keep = np.ones(len(rows), bool)
    for _ in range(4):
        lam = 1e3 * np.sqrt(keep.sum())                        # stiff constraint sum(w) = 1
        w, _ = nnls(np.vstack([A[keep], lam * np.ones((1, det.sum()))]), np.concatenate([np.zeros(keep.sum()), [lam]]),
                    maxiter=50000)
        w = w / w.sum()
        r = A @ w
        sig = 1.4826 * np.median(np.abs(r[keep] - np.median(r[keep])))
        keep = np.abs(r) < 4 * max(sig, 1e-3)
    out = np.zeros(len(jd))
    out[det] = w
    resid = float(np.sqrt(np.mean(r[keep] ** 2)))
    resid_flat = float(np.sqrt(np.mean(A[keep].mean(axis=1) ** 2)))
    return out, int(keep.sum()), resid, resid_flat