"""Per-star detection thresholds for each neutron-capture element in GALAH DR4.

DR4 reports a best-fit [X/Fe] for every star in which the element was fitted and marks it either
'detection' (flag 0; flag 8 = at/above the upper edge of the model range) or 'upper limit'
(flag bit 1: the best-fit model is not distinguishable at 3 sigma from the lowest-abundance model).
We model P(detected | fitted [X/Fe], Teff, logg, [Fe/H], S/N, vsini) with a logistic regression on
the FGK dwarf sample itself and invert it to x90, the abundance at which a star with these
parameters would be reported as a detection 90% of the time.  Stars where the element was not
fitted at all (flag bit 2) get x90 = +inf (not searched).
"""
import numpy as np
from common import *

SNR_COL = {'rb': 'snr_px_ccd4', 'sr': 'snr_px_ccd3', 'y': 'snr_px_ccd1', 'zr': 'snr_px_ccd1',
           'mo': 'snr_px_ccd2', 'ru': 'snr_px_ccd1', 'ba': 'snr_px_ccd2', 'la': 'snr_px_ccd1',
           'ce': 'snr_px_ccd1', 'nd': 'snr_px_ccd2', 'sm': 'snr_px_ccd1', 'eu': 'snr_px_ccd2'}


def design(x, teff, logg, feh, lsnr, lvs):
    t = (teff - 5500.0) / 500.0
    return np.column_stack([np.ones_like(x), x, feh, lsnr, t, t * t, logg - 4.2, lvs,
                            x * t, feh * feh])


def fit_logistic(A, y, n_iter=30, ridge=1e-3):
    b = np.zeros(A.shape[1])
    for _ in range(n_iter):
        z = np.clip(A @ b, -30, 30)
        p = 1.0 / (1.0 + np.exp(-z))
        W = p * (1 - p) + 1e-9
        H = (A * W[:, None]).T @ A + ridge * np.eye(A.shape[1])
        g = A.T @ (y - p) - ridge * b
        step = np.linalg.solve(H, g)
        b = b + step
        if np.max(np.abs(step)) < 1e-6:
            break
    return b


def main(verbose=True):
    d = np.load(CACHE + 'dr4.npz')
    teff = d['teff'].astype(float); logg = d['logg'].astype(float); feh = d['fe_h'].astype(float)
    vs = d['vsini'].astype(float)
    lvs = np.log10(np.clip(np.where(np.isfinite(vs), vs, 5.0), 1.0, 100.0))
    base = (logg >= LOGG_MIN) & (teff < TEFF_MAX) & (teff >= TEFF_MIN) & np.isfinite(feh)
    out = {}; coefs = {}; calib = []
    for e in NCAP:
        x = d[e + '_fe'].astype(float); f = d['flag_' + e + '_fe']
        snr = d[SNR_COL[e]].astype(float)
        lsnr = np.log10(np.clip(np.where(np.isfinite(snr), snr, 1.0), 1.0, 1e4))
        fitted = base & ((f & 2) == 0) & ((f & 4) == 0) & ((f & 64) == 0) & np.isfinite(x) & np.isfinite(snr)
        det = (f & 1) == 0
        A = design(x[fitted], teff[fitted], logg[fitted], feh[fitted], lsnr[fitted], lvs[fitted])
        y = det[fitted].astype(float)
        b = fit_logistic(A, y)
        coefs[e] = b
        # invert for x at P=0.9 and P=0.5 for every dwarf:  z = c0 + x*(b1 + b8*t) + rest
        t = (teff - 5500.0) / 500.0
        rest = (b[0] + b[2] * feh + b[3] * lsnr + b[4] * t + b[5] * t * t + b[6] * (logg - 4.2)
                + b[7] * lvs + b[9] * feh * feh)
        slope = b[1] + b[8] * t
        x90 = np.full(len(x), np.inf); x50 = np.full(len(x), np.inf)
        ok = base & ((f & 2) == 0) & np.isfinite(snr) & (slope > 0.5)
        x90[ok] = (np.log(9.0) - rest[ok]) / slope[ok]
        x50[ok] = (0.0 - rest[ok]) / slope[ok]
        out['x90_' + e] = x90.astype('f4'); out['x50_' + e] = x50.astype('f4')
        # calibration: observed detection fraction vs predicted
        p = 1 / (1 + np.exp(-np.clip(A @ b, -30, 30)))
        for lo, hi in [(0, .1), (.1, .3), (.3, .5), (.5, .7), (.7, .9), (.9, .97), (.97, 1.0001)]:
            m = (p >= lo) & (p < hi)
            calib.append((e, lo, hi, int(m.sum()), float(y[m].mean()) if m.sum() else np.nan,
                          float(p[m].mean()) if m.sum() else np.nan))
        # direct check: among fitted stars with x >= x90, what fraction is detected
        mm = fitted & (x >= x90)
        if verbose:
            print(f'{e}: fitted {fitted.sum()}, detected {int(y.sum())}, coef x={b[1]:.2f} feh={b[2]:.2f} '
                  f'lsnr={b[3]:.2f}; detected fraction among x>=x90: {det[mm].mean() if mm.sum() else np.nan:.3f} (n={mm.sum()}); '
                  f'N dwarfs with x90<=0.5: {(base&(x90<=0.5)).sum()}, <=1.0: {(base&(x90<=1.0)).sum()}')
    np.savez(CACHE + 'dr4_x90.npz', **out)
    import csv
    with open(DATA + 'detection_model_calibration.csv', 'w', newline='') as fh:
        w = csv.writer(fh); w.writerow(['element', 'p_lo', 'p_hi', 'n', 'observed_detected_fraction', 'mean_predicted'])
        for r in calib:
            w.writerow([r[0], r[1], r[2], r[3], f'{r[4]:.4f}', f'{r[5]:.4f}'])
    with open(DATA + 'detection_model_coefficients.csv', 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['element', 'const', 'x', 'feh', 'log10snr', 't', 't2', 'logg-4.2', 'log10vsini', 'x*t', 'feh2'])
        for e in NCAP:
            w.writerow([e] + [f'{v:.5f}' for v in coefs[e]])


if __name__ == '__main__':
    main()