"""Figures for the stellar_salting study."""
import json, csv
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from common import *
import s03_candidates as s3, s04_vetting as s4, s05_physics as s5, spectra

EL_LABEL = {'nd': 'Nd', 'la': 'La', 'ce': 'Ce', 'sm': 'Sm', 'eu': 'Eu', 'zr': 'Zr', 'ba': 'Ba', 'y': 'Y'}
COL = {'nd': '#1b6ca8', 'la': '#d1495b', 'ce': '#edae49', 'sm': '#66a182', 'eu': '#8d5a97', 'zr': '#00798c',
       'ba': '#555555', 'y': '#a0a0a0'}
plt.rcParams.update({'font.size': 10, 'axes.spines.top': False, 'axes.spines.right': False, 'figure.dpi': 110})


def fig_sensitivity(d, c):
    lim = pd.read_csv(DATA + 'limits_by_element.csv')
    nt = pd.read_csv(DATA + 'neff_by_teff.csv')
    fig, ax = plt.subplots(1, 3, figsize=(16.5, 4.3))
    a = ax[2]
    try:
        inj = pd.read_csv(DATA + 'injection_recovery.csv'); inj = inj[inj.salt_added_dex > 0]
        for e in ['nd', 'la']:
            g = inj[inj.element == e]
            a.errorbar(g.salt_added_dex, g.eff_refute_at_S, yerr=[g.eff_refute_at_S - g.eff_refute_at_S_lower95, 0 * g.eff_refute_at_S],
                       fmt='o-', color=COL[e], lw=1.8, ms=5, capsize=3, label='%s: not refuted at the salt level' % EL_LABEL[e])
            a.plot(g.salt_added_dex, g.frac_supported_3sigma, 's:', color=COL[e], lw=1.2, ms=4, mfc='none',
                   label='%s: recovered at >= 3 sigma' % EL_LABEL[e])
        key = json.load(open(DATA + 'key_numbers.json'))
        a.axhline(key['twin_efficiency_used'], color='k', lw=0.8, ls='--')
        a.text(0.76, key['twin_efficiency_used'] - 0.012, 'vetting efficiency used in the limits', fontsize=8, va='top')
        a.set_ylim(0.8, 1.01); a.set_xlabel('salt injected (dex added to the natural prediction)')
        a.set_ylabel('fraction of injected stars'); a.grid(alpha=0.25); a.legend(fontsize=7.5, frameon=False, loc='lower right')
        a.set_title('Injection-recovery of the spectrum test\n(250 random stars of the effective sample per element)', fontsize=10)
    except Exception as ex:
        a.text(0.5, 0.5, 'injection table missing', ha='center', transform=a.transAxes)
    a = ax[0]
    for e in s4.MAIN:
        l = lim[lim.element == e]
        a.plot(l.salt_level_dex, l.N_eff, 'o-', color=COL[e], label=EL_LABEL[e], lw=1.8, ms=4)
        a.plot(l.salt_level_dex, l.N_eff_no_extrapolation, ':', color=COL[e], lw=1.2)
    a.set_yscale('log'); a.set_xlabel('salting level (dex above the natural pattern)')
    a.set_ylabel('effective number of FGK dwarfs/subgiants searched')
    a.set_title('GALAH DR4: stars that would have shown a salt of this size\n(dotted: no credit beyond the 1 dex level)', fontsize=10)
    a.legend(ncol=4, fontsize=8, frameon=False, loc='lower right'); a.grid(alpha=0.25)
    a = ax[1]
    sub = nt[(nt.T_excess == 0.75)]
    tb = sorted(sub.teff_lo.unique())
    for e in ['nd', 'la', 'ce', 'sm', 'eu', 'zr']:
        s = sub[sub.element == e].sort_values('teff_lo')
        a.plot(s.teff_lo + 125, s.N_eff / s.n_clean_stars, 'o-', color=COL[e], label=EL_LABEL[e], lw=1.8, ms=4)
    a.set_xlabel('Teff (K)'); a.set_ylabel('fraction of clean stars sensitive to a 1.0 dex salt')
    a.set_title('Sensitivity by temperature (salt level 1.0 dex)')
    a.set_ylim(0, 1.02); a.legend(ncol=3, fontsize=8, frameon=False); a.grid(alpha=0.25)
    fig.tight_layout(); fig.savefig(FIG + 'fig1_effective_sample.png', dpi=150); plt.close(fig)


def fig_pattern(d, c):
    vet = pd.read_csv(DATA + 'vetting_twin_test.csv')
    surv = pd.read_csv(DATA + 'surviving_candidates_all.csv')
    clean = c['base'] & (d['flag_sp'] == 0)
    fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.6))
    for a, e in zip(ax, ['nd', 'la']):
        m = clean & c['sec'][e] & c['sec']['ba']
        x = d['ba_fe'][m].astype(float); y = d[e + '_fe'][m].astype(float)
        hb = a.hexbin(x, y, gridsize=70, extent=(-0.8, 1.6, -0.8, 1.6), bins='log', cmap='Greys', mincnt=1, lw=0)
        # natural mixing lines
        S = S_FRAC; fs = np.logspace(-0.5, 2, 100)
        for fr, ls, lab in [(1.0, '-', 's-process added to solar mix'), (10 ** 0.4, '--', 'same, [Eu/Fe]=+0.4')]:
            bx = np.log10(S['ba'] * fs + (1 - S['ba']) * fr); by = np.log10(S[e] * fs + (1 - S[e]) * fr)
            a.plot(bx, by, color='#2a9d8f', ls=ls, lw=1.5, label=lab)
        frr = np.logspace(0, 2.2, 50)
        a.plot(np.log10(S['ba'] + (1 - S['ba']) * frr), np.log10(S[e] + (1 - S[e]) * frr), color='#8d5a97', ls=':', lw=1.5,
               label='r-process added to solar mix')
        # fission track (U-235 waste added to a solar-pattern star)
        E, chain = s5.chain_yields('U_235'); ey = s5.elemental_yields(chain)
        tx = []; ty = []
        for dx in np.linspace(0.0, 2.0, 41):
            p = s5.fission_pattern(ey, 'nd', max(dx, 1e-3))[0]
            tx.append(p['ba']); ty.append(p[e])
        a.plot(tx, ty, color='#e76f51', lw=2, label='U-235 fission waste added')
        a.plot([0, 0], [0, 1.5], color='#1b6ca8', lw=2, label='single-element salt')
        v = vet[(vet.kind == 'candidate') & (vet.element == e)]
        i = v.idx.values
        a.scatter(d['ba_fe'][i], d[e + '_fe'][i], s=10, facecolors='none', edgecolors='#d1495b', lw=0.6,
                  label='catalogue candidates (excess >= 0.75 dex)')
        sv = surv[(surv.element == e) & (surv.verdict == 'supported') & (surv.max_level_T >= 0.75)]
        if len(sv):
            a.scatter(sv.ba_fe, sv[e + '_fe'], s=45, marker='*', color='k', zorder=5,
                      label='excess confirmed in the spectrum')
        a.set_xlim(-0.8, 1.6); a.set_ylim(-0.8, 1.6)
        a.set_xlabel('[Ba/Fe]'); a.set_ylabel('[%s/Fe]' % EL_LABEL[e])
        a.set_title('%s versus Ba: natural locus and outliers' % EL_LABEL[e]); a.grid(alpha=0.2)
    ax[1].legend(fontsize=7, frameon=True, framealpha=0.9, loc='lower right')
    fig.tight_layout(); fig.savefig(FIG + 'fig2_natural_pattern.png', dpi=150); plt.close(fig)


def fig_twin(d, c):
    vet = pd.read_csv(DATA + 'vetting_twin_test.csv')
    masks = s4.line_masks()
    fig = plt.figure(figsize=(12, 7.2))
    gs = fig.add_gridspec(2, 4)
    a = fig.add_subplot(gs[0, :2])
    bins = np.linspace(-1.5, 2.5, 41)
    ok = (vet.status == 'ok') & vet.rho.notna() & (vet.snr_mod >= 3)
    cnd = vet[(vet.kind == 'candidate') & ok]
    pos = vet[(vet.kind == 'positive_control') & ok]
    a.hist(np.clip(cnd.rho, -1.5, 2.5), bins=bins, color='#d1495b', alpha=0.75,
           label='catalogue candidates, excess >= 0.75 dex (n=%d)' % len(cnd))
    a.hist(np.clip(pos.rho, -1.5, 2.5), bins=bins, color='#2a9d8f', alpha=0.8,
           label='s-process-rich dwarfs, positive controls (n=%d)' % len(pos))
    a.axvline(0, color='k', lw=0.6); a.axvline(1, color='k', lw=0.6, ls='--')
    a.set_xlabel('rho = observed extra absorption / extra absorption implied by the catalogue abundance')
    a.set_ylabel('star-element pairs'); a.set_yscale('log'); a.legend(fontsize=8, frameon=False)
    a.set_title('Twin test: are the catalogue abundances in the spectra?')
    a = fig.add_subplot(gs[0, 2:])
    zb = np.linspace(-6, 12, 37)
    for kind, col, lab in [('null_control', '#888888', 'ordinary stars'), ('dr3_only', '#edae49', 'DR3 says enhanced, DR4 does not'),
                           ('dr3_both', '#2a9d8f', 'DR3 and DR4 both say enhanced')]:
        g = vet[(vet.kind == kind) & (vet.status == 'ok') & (vet.element == 'nd')]
        if len(g):
            a.hist(np.clip(g.z_null, -6, 12), bins=zb, histtype='step', lw=1.8, color=col, density=True, label='%s (n=%d)' % (lab, len(g)))
    xx = np.linspace(-6, 6, 200); a.plot(xx, np.exp(-xx ** 2 / 2) / np.sqrt(2 * np.pi), 'k:', lw=1)
    a.set_xlabel('extra absorption in the Nd II windows relative to twins (sigma)'); a.set_ylabel('normalised counts')
    a.legend(fontsize=8, frameon=False); a.set_title('Completeness control: does DR4 miss real Nd?')
    # spectra: a refuted Nd catalogue candidate and a positive control, two Nd II lines each
    def panel(ax_list, sid, title):
        i = np.where(d['sobject_id'] == sid)[0][0]
        S = spectra.read(sid); tw = s4.pick_twins(d, c, [i])[i]
        T = [spectra.read(d['sobject_id'][j]) for j in tw]
        T = [t for t in T if t is not None]
        for axx, w0 in zip(ax_list, [4811.342, 5740.858]):
            k = (S['wave'] > w0 - 1.5) & (S['wave'] < w0 + 1.5)
            ref = np.nanmedian([np.interp(S['wave'][k], t['wave'], t['sob']) for t in T], axis=0)
            axx.plot(S['wave'][k], ref, color='#2a9d8f', lw=1.6, label='median of twin stars (observed)')
            axx.plot(S['wave'][k], S['sob'][k], 'k-', lw=1.3, label='star (observed)')
            axx.plot(S['wave'][k], S['smod'][k], color='#d1495b', lw=1.3, ls='--', label='pipeline best-fit model')
            m = [mm for mm in masks['nd'] if abs(mm[0] - w0) < 0.01][0]
            axx.axvspan(m[1], m[2], color='#edae49', alpha=0.25, lw=0)
            axx.set_ylim(0.55, 1.1); axx.set_xlabel('wavelength (A)')
            axx.set_title('%s\nNd II %.1f A' % (title, w0), fontsize=8)
        return i
    cn = vet[(vet.kind == 'candidate') & (vet.element == 'nd') & vet.rho.notna() & (vet.snr_mod >= 6)].sort_values('ex_cat', ascending=False)
    sid_c = int(cn.iloc[0].sobject_id) if len(cn) else None
    pc = vet[(vet.kind == 'positive_control') & (vet.element == 'nd') & vet.rho.notna()].sort_values('snr_mod', ascending=False)
    sid_p = int(pc.iloc[0].sobject_id) if len(pc) else None
    axs = [fig.add_subplot(gs[1, k]) for k in range(4)]
    if sid_c:
        panel(axs[:2], sid_c, 'candidate: catalogue [Nd/Fe]>=%.1f, [Ba/Fe]=%.2f' % (cn.iloc[0].x_cat, d['ba_fe'][int(cn.iloc[0].idx)]))
    if sid_p:
        panel(axs[2:], sid_p, 's-rich control: [Nd/Fe]=%.2f, [Ba/Fe]=%.2f' % (pc.iloc[0].x_cat, d['ba_fe'][int(pc.iloc[0].idx)]))
    axs[0].legend(fontsize=6.5, frameon=False, loc='lower left'); axs[0].set_ylabel('normalised flux')
    fig.tight_layout(); fig.savefig(FIG + 'fig3_twin_test.png', dpi=150); plt.close(fig)


def fig_cost(d=None, c=None):
    """Left: price of a +1 dex salt versus Teff.  Right (revision): what the flag_sp == 0 cut hides --
    a lanthanide-rich 'line forest' star among the flagged dwarfs, compared with ordinary twins."""
    t = pd.read_csv(DATA + 'convective_envelope_mass.csv')
    fig = plt.figure(figsize=(13.5, 4.8))
    gs = fig.add_gridspec(2, 5)
    a = fig.add_subplot(gs[:, :2])
    ms = t[(t.logg >= 4.0)]
    for age, mk, al in [(1.0, 'o', 0.45), (2.0, 's', 0.7), (4.57, '^', 1.0)]:
        s_ = ms[np.abs(ms.age_gyr - age) < 0.35].sort_values('mass_msun')
        a.plot(s_.teff_K, s_['kg_Nd_for_+1dex'], mk + '-', label='Nd, age %.1f Gyr' % age, color='#1b6ca8', alpha=al)
    s_ = ms[np.abs(ms.age_gyr - 2.0) < 0.35].sort_values('mass_msun')
    a.plot(s_.teff_K, s_['kg_Pr_for_+1dex'], 's--', color='#d1495b', label='Pr, age 2.0 Gyr')
    a.set_yscale('log'); a.set_xlabel('Teff (K) of solar-metallicity main-sequence model (0.6-1.5 Msun)'); a.set_ylabel('mass needed for +1.0 dex (kg)')
    a.axhline(9.4e20, color='k', lw=0.6, ls=':'); a.text(4050, 1.1e21, 'mass of Ceres', fontsize=8)
    a.axhline(5e18, color='k', lw=0.6, ls=':'); a.text(4050, 5.8e18, "all Nd in Earth's silicate mantle + crust", fontsize=8)
    a.axvspan(TEFF_MIN, TEFF_MAX, color='#edae49', alpha=0.15, lw=0); a.text(4600, 3e17, 'this search', fontsize=9)
    a.set_title('Price of a +1 dex salt (dilution in the convective envelope)', fontsize=10); a.legend(fontsize=8, frameon=False, loc='lower left'); a.grid(alpha=0.25)
    # --- right: flagged Ap-like star
    try:
        fs = json.load(open(DATA + 'flagged_summary.json'))
        stars = fs['vetting']['nd'].get('class_A_stars', [])
        mm = pd.read_csv(DATA + 'flagged_supported_multi_element.csv')
        mm = mm[mm.sobject_id.isin(stars) & (mm.element == 'nd')].sort_values('snr_px_ccd2')
        sid = int(mm[mm.vsini < 15].iloc[0].sobject_id) if (mm.vsini < 15).any() else int(mm.iloc[0].sobject_id)
        masks = s4.line_masks()
        i = np.where(d['sobject_id'] == sid)[0][0]
        S = spectra.read_compact(sid, masks); tw = s4.pick_twins(d, c, [i])[i]
        T = [spectra.read_compact(d['sobject_id'][j], masks) for j in tw]; T = [x for x in T if x is not None]
        sel = [('nd', 5811.57, 'Nd II'), ('ce', 4773.941, 'Ce II'), ('eu', 6645.0978, 'Eu II'), ('sr', 6550.244, 'Sr I'),
               ('ba', 5853.668, 'Ba II'), ('la', 5805.77, 'La II (+ Nd II 5804.0)')]
        for k, (el, w0, lab) in enumerate(sel):
            ax = fig.add_subplot(gs[k // 3, 2 + k % 3])
            m = [q for q in masks[el] if abs(q[0] - w0) < 0.01][0]
            kk = (S['wave'] > m[3]) & (S['wave'] < m[4])
            ref = np.nanmedian([np.interp(S['wave'][kk], x['wave'], x['sob']) for x in T], axis=0)
            ax.plot(S['wave'][kk], ref, color='#2a9d8f', lw=1.5, label='median of twin stars')
            ax.plot(S['wave'][kk], S['sob'][kk], 'k-', lw=1.1, label='flagged star')
            ax.axvspan(m[1], m[2], color='#edae49', alpha=0.3, lw=0)
            ax.set_ylim(0.45, 1.12); ax.tick_params(labelsize=7); ax.set_title('%s %.1f A' % (lab, w0), fontsize=8)
            if k == 0:
                ax.legend(fontsize=6.5, frameon=False, loc='lower left')
            if k >= 3:
                ax.set_xlabel('wavelength (A)', fontsize=8)
        fig.text(0.70, 0.985, 'What the quality cut hides: GALAH %d (Teff %.0f K, flag_sp = %d), a lanthanide-rich line-forest star'
                 % (sid, d['teff'][i], d['flag_sp'][i]), ha='center', va='top', fontsize=9.5)
    except Exception as ex:
        fig.text(0.7, 0.5, 'flagged-star spectra unavailable (%s)' % ex, ha='center')
    fig.tight_layout(rect=(0, 0, 1, 0.96)); fig.savefig(FIG + 'fig4_cost_and_flagged_cp_star.png', dpi=150); plt.close(fig)
    import os
    old_ = FIG + 'fig4_cost_and_cp_onset.png'
    if os.path.exists(old_):
        os.remove(old_)      # first-pass panel 'CP onset in GALAH' was withdrawn (referee major point 3)


def fig_limits():
    lim = pd.read_csv(DATA + 'limits_by_element.csv'); key = json.load(open(DATA + 'key_numbers.json'))
    sadd = pd.read_csv(DATA + 'limits_salt_added.csv'); sadd = sadd[sadd.adopted == 1]
    nmw = key['galaxy']['n_fgk_milky_way']
    fig, a = plt.subplots(figsize=(8.2, 5.0))
    for e in s4.MAIN:
        l = lim[lim.element == e]
        v = l[l.vetted]; u_ = l[~l.vetted]
        a.plot(v.salt_level_dex, v.f_ul95, 'o-', color=COL[e], lw=1.8, ms=5, label=EL_LABEL[e])
        a.plot(u_.salt_level_dex, u_.f_ul95_catalogue_only, 'o', color=COL[e], ms=4, mfc='none')
        if e in ('nd', 'la'):
            g = sadd[sadd.element == e]
            lo = v.f_ul95.values; hi = np.maximum(g.f_ul95_no_extrapolation.values, v.f_ul95_no_extrapolation.values)
            a.fill_between(v.salt_level_dex.values, lo, hi, color=COL[e], alpha=0.18, lw=0)
    a.set_yscale('log'); a.set_xlabel('salting level (dex above the natural pattern)')
    a.set_ylabel('95% upper limit on the fraction of clean-spectrum FGK dwarfs salted')
    a2 = a.twinx(); a2.set_yscale('log'); lo, hi = a.get_ylim(); a2.set_ylim(lo * nmw, hi * nmw)
    a2.set_ylabel('equivalent number in the Milky Way (%.1e FGK dwarfs)' % nmw); a2.spines['right'].set_visible(True)
    try:
        pr = json.load(open(DATA + 'pr_direct_search_summary.json'))
        xs = [p for p in (1.25, 1.5, 2.0) if pr.get('f_ul95_pr_%.2f' % p)]
        a.plot(xs, [pr['f_ul95_pr_%.2f' % p] for p in xs], 'D--', color='0.45', ms=4, lw=1.0, mfc='none',
               label='Pr (indicative only)')
        a.set_xlim(0.45, 2.05)
    except Exception:
        pass
    a.legend(ncol=4, fontsize=8, frameon=False, loc='upper right')
    a.set_title('Limits. Filled: spectrum-vetted (excess over the natural prediction); open: catalogue only.\n'
                'Bands (Nd, La): up to the most conservative variant (salt defined relative to the star\'s own\n'
                'abundance, no credit for sensitivity beyond the 1 dex level)', fontsize=8.5)
    a.grid(alpha=0.25)
    fig.tight_layout(); fig.savefig(FIG + 'fig5_limits.png', dpi=150); plt.close(fig)


def main(d=None, c=None):
    if d is None:
        d = s3.load(); c = s3.compute(d)
    fig_sensitivity(d, c); fig_pattern(d, c); fig_twin(d, c); fig_cost(d, c); fig_limits()


if __name__ == '__main__':
    main()