#!/usr/bin/env python3
"""Final table of individual-galaxy limits: (i) conservative (every outlier counted), (ii) vetted
(outliers with an identified mundane cause in the stored visual/automatic vetting are not counted).

The vetted limits exist only for thresholds at or below vet_candidates.CAND_THR, because only those
galaxies were inspected."""
import os

import numpy as np
import pandas as pd

import analysis_individual as A
import fpk_common as F
import vet_candidates as V

MANUAL = os.path.join(F.HERE, "vetting_manual.csv")


def main():
    S = A.load_samples()
    best = pd.read_csv(os.path.join(F.DATA, "limits_individual_best.csv"))
    man = pd.read_csv(MANUAL, dtype={"id": str}) if os.path.exists(MANUAL) else None
    rows = []
    for name, label in (("sdss", "SDSS-MGS own (unclipped)"), ("6df", "6dFGSv-J (unclipped)")):
        b = best[(best["sample"] == label) & (best.variant == "default") & (best.relation == "FP")]
        for _, r in b.iterrows():
            rows.append(dict(sample=label, kind="conservative", f=r.f, kappa=r.kappa, threshold_dex=r.threshold_dex,
                             N=r.N, N_eff=r.N_eff, M_eff=r.M_eff, n_obs=r.n_obs, q95=r.q95_best))
        if man is None:
            continue
        df, cfg = S[name]
        res = pd.read_csv(F.resid_path(name))
        cand = pd.read_csv(os.path.join(F.DATA, "candidates_%s_auto.csv" % name), dtype={"id": str})
        m = man[man["sample"] == name]
        cand = cand.merge(m[["id", "explained", "class", "note"]], on="id", how="left")
        if cand.explained.isna().any():
            print("WARNING: %d %s candidates have no stored vetting label; they are counted as unexplained"
                  % (cand.explained.isna().sum(), name))
        cand["explained"] = cand.explained.fillna(0).astype(int)
        cand.to_csv(os.path.join(F.DATA, "candidates_%s_vetted.csv" % name), index=False, float_format="%.6g")
        unexpl = cand.dI.values[cand.explained.values == 0]
        info = dict(coef=None)
        fits = pd.read_json(os.path.join(F.DATA, "fp_fits.json"))
        cs = fits[name]["fpI_coef"][1]
        cj = fits[name]["fj_coef"][1]
        allr = []
        for k in A.KAPPAS:
            t = A.limits(df, res[["dI", "dFJ"]], cs, cj, cfg, label, kappa=k,
                         n_obs_override=lambda thr: int(np.sum(unexpl <= thr)))
            t = t[(t.relation == "FP") & (t.threshold_dex <= V.CAND_THR + 1e-9)]
            allr.append(t)
        allr = pd.concat(allr, ignore_index=True)
        allr["_q"] = allr.q_bonf.fillna(np.inf)
        for f, g in allr.groupby("f"):
            r = g.loc[g._q.idxmin()]
            rows.append(dict(sample=label, kind="vetted", f=f, kappa=r.kappa, threshold_dex=r.threshold_dex, N=r.N,
                             N_eff=r.N_eff, M_eff=r.M_eff, n_obs=r.n_obs, q95=r.q_bonf))
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(F.DATA, "limits_final.csv"), index=False, float_format="%.6g")
    pd.set_option("display.width", 200)
    print(out[out.f.isin([0.5, 0.6, 0.75, 0.9, 0.95])].to_string(index=False))
    return out


if __name__ == "__main__":
    main()