"""Direct search for praseodymium (not in any survey catalogue) in the GALAH DR4 spectra on disk.

Pr II lines in the HERMES windows are taken from the Kurucz line list (gf5901.all).  For every FGK
dwarf/subgiant spectrum already downloaded for the twin test we measure, in +-0.20 A windows around
each line, the extra absorption relative to the pipeline's own best-fit model (which carries Pr at
the scaled-solar level at most), remove the population median in Teff bins (model systematics) and
combine the lines with a matched filter.  The stack of s-process-rich dwarfs (which must contain
extra Pr) validates the line list and calibrates the response empirically.
"""
import os, csv, glob, json, urllib.request
import numpy as np
import pandas as pd
from common import *
import s03_candidates as s3
import spectra

KURUCZ = 'http://kurucz.harvard.edu/linelists/gfall/gf5901.all'
WINDOWS = [(4715.0, 4901.5), (5650.0, 5872.0), (6478.5, 6735.0)]
HALF = 0.20      # A, half-width of the line window (Pr II lines are broadened by hyperfine structure)
CTRL_OFF = 0.9   # A, offset of the control windows
Z_SENS = 10.0    # expected significance that defines a star's quoted sensitivity
THETA = 5040.0 / 5500.0


def pr_lines():
    p = RAW + 'lit/gf5901.all'
    os.makedirs(RAW + 'lit', exist_ok=True)
    if not os.path.exists(p):
        with urllib.request.urlopen(KURUCZ, timeout=120) as r:
            open(p, 'wb').write(r.read())
    L = []
    for l in open(p, errors='ignore'):
        try:
            wl = float(l[0:11]) * 10.0; gf = float(l[11:18]); e1 = abs(float(l[24:36])); e2 = abs(float(l[52:64]))
        except ValueError:
            continue
        if l[18:24].strip() != '59.01':
            continue
        chi = min(e1, e2) / 8065.544
        if any(a <= wl <= b for a, b in WINDOWS):
            L.append((wl, gf, chi, gf - THETA * chi))
    L = sorted(L, key=lambda x: -x[3])
    # merge lines closer than the window (blended Pr II pairs): keep the stronger as the centre
    out = []
    for x in L:
        if x[3] < L[0][3] - 0.55:
            break
        if all(abs(x[0] - y[0]) > 2 * HALF for y in out):
            out.append(x)
    return out


def measure(S, lines):
    w = S['wave']; dl = np.gradient(w)
    res = (S['smod'] - S['sob'])           # positive = more absorption than the model
    good = np.isfinite(res) & np.isfinite(S['uob']) & (S['uob'] > 0) & (S['uob'] < 0.3)
    R = np.full(len(lines), np.nan); E = np.full(len(lines), np.nan)
    for k, (wl, gf, chi, st) in enumerate(lines):
        inm = (np.abs(w - wl) <= HALF) & good
        side = (np.abs(w - wl) > 0.45) & (np.abs(w - wl) <= 1.6) & good
        if inm.sum() < 5 or side.sum() < 15:
            continue
        off = np.median(res[side])
        R[k] = np.sum((res[inm] - off) * dl[inm]) * 1e3
        E[k] = np.sqrt(np.sum((S['uob'][inm] * dl[inm]) ** 2)) * 1e3
    return R, E


def reduce(R, E, teff, ref_ok):
    """Remove the population median per line in 250 K bins (model systematics) and add the excess
    scatter of ordinary stars in quadrature to the formal error."""
    Rc = np.full_like(R, np.nan); Et = np.full_like(R, np.nan)
    for t0 in np.arange(TEFF_MIN, TEFF_MAX, 250.0):
        m = (teff >= t0) & (teff < t0 + 250)
        ref = m & ref_ok
        if ref.sum() < 20:
            continue
        for k in range(R.shape[1]):
            v = R[ref, k]; v = v[np.isfinite(v)]
            if len(v) < 20:
                continue
            med = np.median(v); mad = 1.4826 * np.median(np.abs(v - med))
            ef = np.nanmedian(E[ref, k])
            sys2 = max(mad ** 2 - ef ** 2, 0.0)
            Rc[m, k] = R[m, k] - med; Et[m, k] = np.sqrt(E[m, k] ** 2 + sys2)
    return Rc, Et


def combine(Rc, Et, g, use):
    ok = np.isfinite(Rc) & np.isfinite(Et) & use[None, :]
    num = np.where(ok, g[None, :] * Rc / Et ** 2, 0).sum(1); den = np.where(ok, g[None, :] ** 2 / Et ** 2, 0).sum(1)
    W = np.where(den > 0, num / np.where(den > 0, den, 1), np.nan); sW = np.where(den > 0, 1 / np.sqrt(np.where(den > 0, den, 1)), np.nan)
    nsig = (np.where(ok, Rc / Et, 0) > 2).sum(1)
    return W, sW, ok.sum(1), nsig


def main(verbose=True):
    d = s3.load(); c = s3.compute(d)
    lines = pr_lines()
    ctrl = [[(x[0] + off, x[1], x[2], x[3]) for x in lines] for off in (-CTRL_OFF, CTRL_OFF)]
    have = set(int(os.path.basename(f).split('_')[0]) for f in glob.glob(RAW + 'spec/*_allstar_fit_spectrum.fits'))
    clean = c['base'] & (d['flag_sp'] == 0)
    idx = np.array([i for i in np.where(clean)[0] if int(d['sobject_id'][i]) in have])
    nL = len(lines)
    cp = CACHE + 'pr_measurements.npz'      # cache of the per-line measurements (re-made if the spectrum set changes)
    done = False
    if os.path.exists(cp):
        z_ = np.load(cp)
        if len(z_['sid']) == len(idx) and np.array_equal(z_['sid'], d['sobject_id'][idx]) and z_['R'].shape[1] == nL:
            R, E, RC, EC = z_['R'], z_['E'], list(z_['RC']), list(z_['EC']); done = True
    if not done:
        R = np.full((len(idx), nL), np.nan); E = np.full_like(R, np.nan)
        RC = [np.full_like(R, np.nan) for _ in ctrl]; EC = [np.full_like(R, np.nan) for _ in ctrl]
        for n, i in enumerate(idx):
            try:
                S = spectra.read(d['sobject_id'][i])
            except Exception:
                S = None
            if S is None:
                continue
            R[n], E[n] = measure(S, lines)
            for q, cl in enumerate(ctrl):
                RC[q][n], EC[q][n] = measure(S, cl)
        np.savez(cp, sid=d['sobject_id'][idx], R=R, E=E, RC=np.array(RC), EC=np.array(EC))
    teff = d['teff'][idx].astype(float); feh = d['fe_h'][idx].astype(float)
    # s-process-rich stars (they must contain extra Pr) versus ordinary stars
    srich = c['sec']['ba'][idx] & (d['ba_fe'][idx] >= 0.6) & (((d['nd_fe'][idx] >= 0.6) & c['sec']['nd'][idx]) | ((d['la_fe'][idx] >= 0.6) & c['sec']['la'][idx]))
    Rc, Et = reduce(R, E, teff, ~srich)
    g = 10 ** (np.array([x[3] for x in lines]) - lines[0][3])      # relative line strengths (gf, excitation)
    # per-line response of the s-rich stack; a line is kept if its response per unit strength is
    # within a factor 2.5 of that of the strongest line, i.e. it behaves like a Pr line (lines that
    # respond much more are blended with lines of other s-process elements, lines that respond
    # negatively sit on features the pipeline model over-predicts in s-rich stars)
    resp = np.full(nL, np.nan); rerr = np.full(nL, np.nan)
    for k in range(nL):
        mk = srich & np.isfinite(Rc[:, k])
        if mk.sum() > 5:
            resp[k] = np.sum(Rc[mk, k] / Et[mk, k] ** 2) / np.sum(1 / Et[mk, k] ** 2); rerr[k] = 1 / np.sqrt(np.sum(1 / Et[mk, k] ** 2))
    ratio = (resp / g) / (resp[0] / g[0])
    use = np.isfinite(ratio) & (ratio >= 0.4) & (ratio <= 2.5)
    with open(DATA + 'pr_lines_used.csv', 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['wavelength_A', 'loggf_kurucz', 'chi_eV', 'relative_strength', 'srich_stack_mA', 'srich_stack_err', 'response_ratio', 'used'])
        for k, x in enumerate(lines):
            w.writerow(['%.3f' % x[0], '%.2f' % x[1], '%.3f' % x[2], '%.2f' % g[k], '%.2f' % resp[k], '%.2f' % rerr[k], '%.2f' % ratio[k], int(use[k])])
    W, sW, nl, nsig = combine(Rc, Et, g, use)
    zz = (W / sW)[~srich & np.isfinite(W)]
    zscale = max(1.0, 1.4826 * np.median(np.abs(zz - np.median(zz))))
    sW = sW * zscale; z = W / sW
    # control windows shifted by -/+ CTRL_OFF: same pipeline, no Pr line
    zc = []
    for q in range(len(ctrl)):
        Rq, Eq = reduce(RC[q], EC[q], teff, ~srich)
        Wq, sq, _, _ = combine(Rq, Eq, g, use)
        zq = Wq / sq
        zq = zq / max(1.0, 1.4826 * np.nanmedian(np.abs(zq[~srich] - np.nanmedian(zq[~srich]))))
        zc.append(zq)
    zctrl = np.nanmax(np.abs(np.array(zc)), axis=0)
    # --- calibration on s-rich stars.  Their natural [Pr/Fe] is predicted from Ba (solar s-fractions);
    # the combined Pr signal is modelled as
    #     W = k0 * 10^(a [Fe/H] + b (Teff-5500)/1000 + c (logg-4)) * (10^[Pr/Fe] - 1)
    # (weak lines of the dominant ionisation stage), fitted by robust least squares.
    from scipy.optimize import least_squares
    logg = d['logg'][idx].astype(float)
    fr = 10 ** np.clip(-0.4 * feh, 0, 0.4)
    fs = np.maximum((10 ** d['ba_fe'][idx].astype(float) - (1 - S_FRAC['ba']) * fr) / S_FRAC['ba'], 0.05)
    pr_pred = np.log10(S_FRAC['pr'] * fs + (1 - S_FRAC['pr']) * fr)       # natural [Pr/Fe] expected from Ba
    usable = np.isfinite(W) & (nl >= max(3, int(0.6 * use.sum())))
    sr = srich & usable

    def resp_star(p):
        return p[0] * 10 ** (p[1] * feh + p[2] * (teff - 5500.0) / 1000.0 + p[3] * (logg - 4.0))
    fit = least_squares(lambda p: ((W - resp_star(p) * (10 ** pr_pred - 1.0)) / sW)[sr], [1.0, 0.0, 0.0, 0.0], loss='soft_l1')
    pcal = fit.x
    try:
        pcal_err = np.sqrt(np.diag(np.linalg.inv(fit.jac.T @ fit.jac)))
    except Exception:
        pcal_err = np.full(4, np.nan)
    kstar = resp_star(pcal)
    zexp = kstar * (10 ** pr_pred - 1.0) / sW
    stack_mean = np.sum(W[sr] / sW[sr] ** 2) / np.sum(1 / sW[sr] ** 2); stack_err = 1 / np.sqrt(np.sum(1 / sW[sr] ** 2))
    nr = ~srich & usable
    norm_mean = np.sum(W[nr] / sW[nr] ** 2) / np.sum(1 / sW[nr] ** 2); norm_err = 1 / np.sqrt(np.sum(1 / sW[nr] ** 2))
    good = sr & (z > 3) & (W > 0)
    scat = float(np.std(np.log10(W[good] / (kstar[good] * (10 ** pr_pred[good] - 1.0))))) if good.sum() > 5 else np.nan
    # --- per-star sensitivity: the [Pr/Fe] at which the EXPECTED significance is Z_SENS (= 10): with the
    # 0.2 dex scatter of the calibration a star at that level is detected at z >= 5 with ~90% probability
    p5 = np.full(len(idx), np.inf)
    p5[usable] = np.log10(1.0 + Z_SENS * sW[usable] / kstar[usable])
    pr_est = np.full(len(idx), np.nan)
    pos = usable & (W > 0)
    pr_est[pos] = np.log10(1.0 + W[pos] / kstar[pos])
    detect = usable & (z >= 5) & (nsig >= 3)
    cand = detect & ~srich & (zctrl < 3)
    out = pd.DataFrame(dict(sobject_id=d['sobject_id'][idx], ra=d['ra'][idx], dec=d['dec'][idx], teff=teff, logg=d['logg'][idx], fe_h=feh,
                            snr_ccd1=d['snr_px_ccd1'][idx], snr_ccd2=d['snr_px_ccd2'][idx], n_lines=nl, W_mA=W,
                            sW_mA=sW, z=z, z_control_max=zctrl, n_lines_gt2sig=nsig, srich=srich.astype(int), pr_fe_sensitivity=p5,
                            pr_fe_estimate=pr_est, pr_pred_from_ba=pr_pred, ba_fe=d['ba_fe'][idx], nd_fe=d['nd_fe'][idx],
                            flag_nd=d['flag_nd_fe'][idx], la_fe=d['la_fe'][idx], flag_la=d['flag_la_fe'][idx],
                            detected=detect.astype(int), candidate_pr_only=cand.astype(int)))
    out.to_csv(DATA + 'pr_direct_search.csv', index=False, float_format='%.3f')
    key = dict(n_lines_considered=nL, n_lines_used=int(use.sum()), lines_used=[float(lines[k][0]) for k in range(nL) if use[k]],
               n_spectra=int(usable.sum()), n_srich=int(sr.sum()),
               srich_stack_mA=float(stack_mean), srich_stack_err=float(stack_err),
               normal_stack_mA=float(norm_mean), normal_stack_err=float(norm_err),
               calibration=dict(k0_mA=float(pcal[0]), a_feh=float(pcal[1]), b_teff_per_kK=float(pcal[2]), c_logg=float(pcal[3]),
                                errors=[float(v) for v in pcal_err], scatter_dex=scat),
               z_sensitivity=Z_SENS, zscale=float(zscale),
               n_detected=int(detect.sum()), n_detected_srich=int((detect & srich).sum()),
               n_detected_not_srich=int((detect & ~srich).sum()),
               n_detected_not_srich_control_fail=int((detect & ~srich & ~(zctrl < 3)).sum()),
               n_candidates_pr_only=int(cand.sum()),
               frac_ordinary_control_z_gt3=float(np.mean(zctrl[nr] >= 3)))
    # validation: s-rich stars with expected significance >= 5 and >= Z_SENS, and how many are detected
    for zt in (5.0, Z_SENS):
        v = sr & (zexp >= zt)
        key['validation_expected_z_ge_%g' % zt] = dict(n=int(v.sum()), n_detected=int((v & detect).sum()))
    # --- revision (referee major point 5): the Pr numbers are INDICATIVE only.
    # (a) efficiency: the only empirical recovery test is on s-rich calibrators with expected
    #     significance >= 5 (none reaches the level Z_SENS = 10 at which the sensitivity is quoted);
    #     the 95% lower bound of that recovery fraction is applied as the efficiency;
    # (b) the quoted sensitivities ([Pr/Fe] >= 1.5 and 2.0) are a linear (weak-line) extrapolation
    #     beyond the calibrators' predicted [Pr/Fe], with no allowance for line saturation;
    # (c) the stars are those whose spectra were downloaded for the twin test (candidates, controls
    #     and their high-S/N twins), not a random sample of FGK dwarfs.
    v5 = key['validation_expected_z_ge_5']
    eff_pr = v5['n_detected'] / max(v5['n'], 1)
    eff_pr_lo = 1.0 - binom_upper(v5['n'] - v5['n_detected'], v5['n']) if v5['n'] else np.nan
    key['recovery_fraction_calibrators'] = float(eff_pr); key['recovery_fraction_lower95'] = float(eff_pr_lo)
    cal = pr_pred[sr]
    key['calibrator_pr_fe_predicted'] = dict(min=float(np.nanmin(cal)), median=float(np.nanmedian(cal)), max=float(np.nanmax(cal)))
    for p in (1.0, 1.25, 1.5, 2.0):
        sens = (~srich) & (p5 <= p)
        n = int(sens.sum())
        # candidates whose Pr estimate reaches this level (within 0.3 dex) count against the limit
        kc = int((cand & (pr_est >= p - 0.3)).sum())
        key['n_sensitive_pr_%.2f' % p] = n; key['n_candidates_at_%.2f' % p] = kc
        key['f_ul95_pr_%.2f_no_efficiency' % p] = float(poisson_upper(kc) / n) if n > 0 else None
        key['f_ul95_pr_%.2f' % p] = float(poisson_upper(kc) / (n * eff_pr_lo)) if n > 0 else None
        key['extrapolation_beyond_calibrators_dex_%.2f' % p] = float(p - np.nanmax(cal))
        if n > 0:
            key['sensitive_sample_%.2f' % p] = dict(
                median_snr_ccd1=float(np.nanmedian(d['snr_px_ccd1'][idx][sens])), median_snr_ccd2=float(np.nanmedian(d['snr_px_ccd2'][idx][sens])),
                median_teff=float(np.nanmedian(teff[sens])))
    clean_all = c['base'] & (d['flag_sp'] == 0)
    key['all_clean_dwarfs_median_snr_ccd1'] = float(np.nanmedian(d['snr_px_ccd1'][clean_all]))
    key['all_clean_dwarfs_median_snr_ccd2'] = float(np.nanmedian(d['snr_px_ccd2'][clean_all]))
    key['status'] = 'indicative only: see comments in s09_pr_search.py'
    json.dump(key, open(DATA + 'pr_direct_search_summary.json', 'w'), indent=1)
    if verbose:
        print(json.dumps(key, indent=1))
        pd.set_option('display.width', 250); pd.set_option('display.max_columns', 40)
        print(out[(out.detected == 1) & (out.srich == 0)].sort_values('z', ascending=False).head(30).to_string())
    return key, out


if __name__ == '__main__':
    main()