"""Part 2c: a direct measurement of whether the OGLE-II DIA catalogue of candidate variables contains
stars that lose most of their light, using an independent list of such stars.

The search of 04_ogle2_search.py covers the 221,801 objects that the 2002 variable finder chose to
catalogue, not all 6.4 million monitored stars.  Whether a star that vanishes would have been
catalogued was inferred, in the first version of this study, from the recovery of microlensing
*brightenings* (174 of 214; Wozniak et al. 2001).  Here it is measured for *dimmings*:

  The OGLE Collection of Variable Stars lists 425,193 eclipsing binaries toward the bulge
  (Soszynski et al. 2016), found by period searches in OGLE-II/III/IV photometry, independently of
  the Wozniak et al. (2002) catalogue.  59,584 of them carry an OGLE-II identifier, i.e. they are
  stars of the OGLE-II template images.  Those with a primary eclipse deeper than 0.75 mag lose
  more than half of their light at every primary minimum.  For each magnitude bin we count the
  fraction that the DIA catalogue contains (position match), and the fraction that it contains
  with its centroid on a reference-image star (cut C2 of the search).

An eclipse is a harder case for the variable finder than an occultation: the star is below half
light for a few percent of the frames, never on consecutive nights, whereas an occulted star is
at zero for five or more consecutive frames.  The measured fractions are therefore lower bounds on
the probability that a totally occulted star of the same magnitude is catalogued and passes C2.

Outputs: data/eb_inclusion.csv, data/eb_inclusion_summary.json
"""
import json
import os

import numpy as np
import pandas as pd
from astropy.coordinates import SkyCoord
from astropy.io import fits
import astropy.units as u

from common import DATA, RAW
import ogle2lib as ol

MAG_EDGES = [0, 15, 16, 17, 17.5, 18]
DEPTH_BINS = [(0.75, 9.0, "primary eclipse deeper than 0.75 mag (below 50% of the light)"),
              (0.40, 0.75, "primary eclipse 0.40-0.75 mag"),
              (0.0, 0.40, "primary eclipse shallower than 0.40 mag")]
MATCH_ARCSEC = 1.0


def read_ocvs():
    d = os.path.join(RAW, "ocvs")
    rows = []
    with open(os.path.join(d, "ident.dat")) as f:
        for line in f:
            o2 = line[83:98].strip()
            if not o2.startswith("BUL_SC"):
                continue
            ra = (int(line[25:27]) + int(line[28:30]) / 60 + float(line[31:36]) / 3600) * 15.0
            dec = (int(line[38:40]) + int(line[41:43]) / 60 + float(line[44:48]) / 3600) * (-1 if line[37] == "-" else 1)
            rows.append((line[:19].strip(), line[21:24].strip(), ra, dec, int(o2.split("_")[1][2:]), line[50:66].strip() != ""))
    ident = pd.DataFrame(rows, columns=["id", "subtype", "ra", "dec", "field", "has_ogle4_id"])
    par = pd.read_csv(os.path.join(d, "ecl.dat"), sep=r"\s+", header=None, na_values="-",
                      names=["id", "I_max", "V_max", "period", "t0", "depth1", "depth2"])
    return ident.merge(par, on="id", how="inner")


def main():
    eb = read_ocvs()
    n_all = len(eb)
    # keep the binaries that carry an OGLE-IV identifier: they were found in OGLE-IV photometry, so the
    # sample cannot have been drawn from the OGLE-II DIA catalogue that is being tested
    eb = eb[np.isfinite(eb.I_max) & np.isfinite(eb.depth1) & eb.has_ogle4_id].copy()
    sep, idrad, sep_shift = np.full(len(eb), np.nan), np.full(len(eb), np.nan), np.full(len(eb), np.nan)
    eb = eb.reset_index(drop=True)
    for i in range(1, 50):
        s = np.where(eb.field.values == i)[0]
        if not len(s):
            continue
        cat = fits.open(os.path.join(ol.DIA, f"bul_sc{i}_cat4.fts"))[1].data
        cc = SkyCoord(np.array(cat["RA"], float) * u.deg, np.array(cat["DEC"], float) * u.deg)
        ce = SkyCoord(eb.ra.values[s] * u.deg, eb.dec.values[s] * u.deg)
        j, d2, _ = ce.match_to_catalog_sky(cc)
        # remove the mean astrometric offset between the two catalogues, measured on the close pairs
        close = d2.arcsec < 1.5
        if close.sum() > 20:
            dra = np.median(((cc.ra.deg[j] - ce.ra.deg) * np.cos(np.radians(ce.dec.deg)))[close])
            ddec = np.median((cc.dec.deg[j] - ce.dec.deg)[close])
            ce = SkyCoord((eb.ra.values[s] + dra / np.cos(np.radians(eb.dec.values[s]))) * u.deg, (eb.dec.values[s] + ddec) * u.deg)
            j, d2, _ = ce.match_to_catalog_sky(cc)
        sep[s] = d2.arcsec
        idrad[s] = np.array(cat["ID_RAD"], float)[j]
        # chance coincidences: the same match after moving every binary 20 arcsec north
        cs = SkyCoord(ce.ra, ce.dec + 20 * u.arcsec)
        sep_shift[s] = cs.match_to_catalog_sky(cc)[1].arcsec
    eb["sep_arcsec"], eb["id_rad_of_match"] = sep, idrad
    eb["in_catalogue"] = eb.sep_arcsec <= MATCH_ARCSEC
    eb["in_catalogue_c2"] = eb.in_catalogue & (eb.id_rad_of_match <= ol.CUTS["idrad_max"])
    rows = []
    for lo, hi, lab in DEPTH_BINS:
        for b in range(len(MAG_EDGES) - 1):
            s = eb[(eb.depth1 >= lo) & (eb.depth1 < hi) & (eb.I_max >= MAG_EDGES[b]) & (eb.I_max < MAG_EDGES[b + 1])]
            n = len(s)
            k1, k2 = int(s.in_catalogue.sum()), int(s.in_catalogue_c2.sum())
            rows.append({"depth_bin": lab, "I_bin": f"{MAG_EDGES[b]}_{MAG_EDGES[b + 1]}", "n_eclipsing_binaries": n,
                         "n_in_dia_catalogue": k1, "frac_in_dia_catalogue": k1 / n if n else np.nan,
                         "n_in_catalogue_and_pass_C2": k2, "frac_in_catalogue_and_pass_C2": k2 / n if n else np.nan,
                         "binomial_err_C2": np.sqrt(k2 * (n - k2) / n ** 3) if n else np.nan,
                         "median_period_d": float(s.period.median()) if n else np.nan})
    tab = pd.DataFrame(rows)
    tab.to_csv(os.path.join(DATA, "eb_inclusion.csv"), index=False, float_format="%.4g")
    pd.set_option("display.width", 250)
    print(tab.to_string(index=False, float_format=lambda x: f"{x:.3g}"))
    deep = eb[(eb.depth1 >= 0.75) & (eb.I_max < 18)]
    longp = deep[deep.period > 10]
    summary = {
        "n_eclipsing_binaries_with_ogle2_id": int(n_all),
        "n_used_with_ogle4_id_and_parameters": int(len(eb)),
        "chance_match_fraction_after_20arcsec_shift": float((sep_shift <= MATCH_ARCSEC).mean()),
        "n_deep_I_lt_18_between_1_and_2_arcsec": int(((deep.sep_arcsec > 1) & (deep.sep_arcsec <= 2)).sum()),
        "n_deep_I_lt_18": int(len(deep)),
        "frac_deep_I_lt_18_in_catalogue": float(deep.in_catalogue.mean()),
        "frac_deep_I_lt_18_in_catalogue_and_pass_C2": float(deep.in_catalogue_c2.mean()),
        "n_deep_I_lt_18_period_gt_10d": int(len(longp)),
        "frac_deep_period_gt_10d_in_catalogue_and_pass_C2": float(longp.in_catalogue_c2.mean()) if len(longp) else None,
        "frac_deep_in_catalogue_and_pass_C2_by_bin": {r.I_bin: r.frac_in_catalogue_and_pass_C2 for r in tab[tab.depth_bin == DEPTH_BINS[0][2]].itertuples()},
        "frac_deep_in_catalogue_by_bin": {r.I_bin: r.frac_in_dia_catalogue for r in tab[tab.depth_bin == DEPTH_BINS[0][2]].itertuples()},
        "n_deep_by_bin": {r.I_bin: int(r.n_eclipsing_binaries) for r in tab[tab.depth_bin == DEPTH_BINS[0][2]].itertuples()},
        "match_radius_arcsec": MATCH_ARCSEC,
        "source": "OGLE Collection of Variable Stars, eclipsing binaries toward the bulge (Soszynski et al. 2016, Acta Astron. 66, 405)",
    }
    with open(os.path.join(DATA, "eb_inclusion_summary.json"), "w") as f:
        json.dump(summary, f, indent=1)
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()