"""Collect compact chart data from the study outputs and the model results into site/data.js (window.FERMI = {...})."""
import json, os
import numpy as np
import pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
S = ROOT + '/studies'
out = {}


def r(x, n=3):
    return [None if (v is None or (isinstance(v, float) and not np.isfinite(v))) else round(float(v), n) for v in x]


# ---- flyby: every Gaia star with a measured radial velocity that passes within 10 pc in +-1 Myr-ish ----
try:
    d = pd.read_csv(S + '/flyby/data/census_encounters_100pc.csv')
    d = d[(d.d_ph <= 10) & (d.t_ph.abs() <= 1.0)]
    out['flyby'] = dict(t=r(d.t_ph, 3), d=r(d.d_ph, 2), v=r(d.v, 0), MG=r(d.MG, 1), n=int(len(d)))
    e = pd.read_csv(S + '/flyby/data/exposure_table.csv')
    out['flyby_exposure'] = e[['era', 'T_Myr', 'd_pc', 'N_central', 'N_low', 'N_high']].round(1).to_dict('records')
except Exception as ex:
    print('flyby failed', ex)

# ---- headstart: formation-age distribution of Sun-like stars at the Sun's radius ----
try:
    d = pd.read_csv(S + '/headstart/data/xr_age_distribution.csv')     # primary sample: LAMOST subgiants, Xiang & Rix 2022 ages
    out['headstart'] = dict(lo=r(d.age_lo, 2), hi=r(d.age_hi, 2), all=r(d['formed: all [Fe/H]'], 5), metal=r(d['formed: [Fe/H]>-0.5'], 5))
except Exception as ex:
    print('headstart failed', ex)

# ---- fundamental plane: residual histogram and limits ----
try:
    d = pd.read_csv(S + '/fp_kiii/data/resid_sdss.csv', usecols=['dI', 'clean'])
    edges = np.round(np.arange(-1.2, 0.8001, 0.025), 3)
    h, _ = np.histogram(d.dI.clip(-1.2, 0.8), edges)
    lim = pd.read_csv(S + '/fp_kiii/data/limits_final.csv')
    out['fp'] = dict(edges=edges.tolist(), counts=h.tolist(), n=int(len(d)),
                     limits=lim[['sample', 'kind', 'f', 'N_eff', 'n_obs', 'q95']].round(5).to_dict('records'))
except Exception as ex:
    print('fp failed', ex)
try:
    lim = pd.read_csv(S + '/tf_kiii/data/limits.csv')
    out['tf'] = lim[lim.case.str.startswith('PRIMARY')][['f_min', 'N_unresolved', 'N_eff', 'UL95_fraction']].round(5).to_dict('records')
except Exception as ex:
    print('tf failed', ex)

# ---- SETI: maximum number of transmitters in the Galaxy vs power ----
try:
    d = pd.read_csv(S + '/seti_meta/data/limits_vs_eirp.csv')
    cols = {'hi': 'Nmax95_MW_HI_line', 'band': 'Nmax95_MW_log_1_10GHz', 'wide': 'Nmax95_MW_log_0p1_100GHz'}
    out['seti'] = dict(eirp=[float(f'{v:.4g}') for v in d.EIRP_W], **{k: [float(f'{v:.4g}') if np.isfinite(v) else None for v in d[c]] for k, c in cols.items() if c in d})
except Exception as ex:
    print('seti failed', ex)

# ---- timing: events and the fast:slow abiogenesis ratio vs number of later hard steps ----
try:
    d = pd.read_csv(S + '/timing/data/events.csv')
    out['events'] = d[['key', 'group', 'event', 'age_Ga', 'lo_Ga', 'hi_Ga', 'status']].round(4).to_dict('records')
    b = pd.read_csv(S + '/timing/data/a1_bmax_table.csv')
    out['bmax'] = b.round(3).to_dict('records')
except Exception as ex:
    print('timing failed', ex)

# ---- earliness: rank vs hard steps and lifetime cap ----
try:
    d = pd.read_csv(S + '/mdwarf_early/data/rank_grid_n_Lmax.csv')
    out['rank'] = dict(n=r(d.n, 0), L=r(d.Lmax_Gyr, 2), all=[float(f'{v:.3g}') for v in d.rank_all_stars], noM=[float(f'{v:.3g}') for v in d.rank_hosts_above_0p6])
    c = pd.read_csv(S + '/mdwarf_early/data/rocky_atmosphere_census.csv')
    out['census'] = c[['planet', 'host_class', 'host_mass_Msun', 'radius_Re', 'S_earth', 'v_esc_kms', 'outcome', 'pi_base']].round(3).to_dict('records')
except Exception as ex:
    print('mdwarf failed', ex)

# ---- convergence: encephalisation through time ----
try:
    d = pd.read_csv(S + '/convergence/data/part2_brain_size_through_time.csv')
    d['eqv'] = pd.to_numeric(d['eq'], errors='coerce'); d['age'] = pd.to_numeric(d['age_Ma'], errors='coerce')
    d = d[np.isfinite(d['eqv']) & np.isfinite(d['age'])]
    out['brains'] = dict(age=r(d['age'], 2), eq=r(d['eqv'], 3), clade=d['clade'].tolist(), species=d['species'].tolist())
except Exception as ex:
    print('convergence failed', ex)

# ---- model (version 4) ----
M = ROOT + '/models'
res = json.load(open(M + '/integrate4_results.json'))
CLS = res['classes']
r3 = lambda v: round(float(v), 4)


def slim_school(q):
    o = dict(posterior={c: r3(q['posterior'][c]) for c in CLS}, alone=r3(q['p_no_other_ever_in_galaxy']), now=r3(q['p_other_now_in_galaxy']),
             anyone_now=r3(q['p_anyone_now_in_galaxy']), alone_exact=r3(q['p_no_other_ever_poisson']), life_common=r3(q['p_life_common']),
             sub={k: r3(v) for k, v in q['sub'].items()}, ess=round(q['ess']), predictive=r3(q.get('predictive', 1.0)),
             p_barrier=r3(q['p_barrier']), p_fast=r3(q['p_fast']), p_gated=r3(q['p_gated']), steps=r3(q['mean_hard_steps']),
             hold1=r3(q['p_expanders_hold_1pct']), claim1=r3(q['p_claimers_hold_1pct']), reach=r3(q['p_nobody_in_reach']), reach_c=r3(q['p_nobody_in_reach_c']),
             reach_exp=r3(q['p_no_expander_in_reach']), regime={k: r3(v) for k, v in q['regime'].items()},
             p_evident=float(f"{q['mean_p_exp_evident']:.3g}"), pexp_lt=r3(q['p_pexp_lt_1e6']), mean_like={c: r3(q['mean_like'][c]) for c in CLS},
             evidence=float(f"{q['evidence']:.5g}"), nq=q['log10_N_ever_q'], lineages=q['log10_lineages_unnoticed_q'])
    if 'predictive_true' in q:
        o['predictive_true'] = r3(q['predictive_true'])
    return o


def slim_run(x):
    return dict(prior={c: r3(x['prior'][c]) for c in CLS}, prior_alone=r3(x['prior_stats']['p_no_other_ever_in_galaxy']),
                prior_sub={k: r3(v) for k, v in x['prior_stats']['sub'].items()},
                schools={s: slim_school(q) for s, q in x['schools'].items()})


# the seven explanations in plain words for the chart tooltips (same meaning as res['class_labels'])
PLAIN_LABELS = dict(
    empty_universe='No other civilisation has ever arisen anywhere in the universe we can observe',
    rare='Others exist elsewhere in the universe, but none has ever arisen in our Galaxy and none is on its way',
    en_route='Expanding civilisations already hold more than 1% of the universe but have not reached us',
    stuck='Civilisations have arisen in our Galaxy, and none has kept spreading from star to star',
    unnoticed='Settlers have spread through the Galaxy but leave so little trace that we have not noticed them',
    hidden='Settlers have spread through the Galaxy and are deliberately hiding',
    excluded='Settlers have visibly claimed the Galaxy and we sit in a pocket they left alone (contradicted by what we see, unless such pockets exist)',
)
assert set(PLAIN_LABELS) == set(res['class_labels'])
b0 = res['base']
out['model'] = dict(
    classes=CLS, class_labels=PLAIN_LABELS, schools=res['schools'], school_labels=res['school_labels'],
    base=dict(slim_run(b0), prior_stats=b0['prior_stats']),
    shapley=res['shapley_alone'], shap=res.get('shapley'), early_factor=res.get('early_life_factor'), version=res.get('model_version', '4'),
    ablation={k: slim_run(v) for k, v in res['ablation'].items()},
    camps={k: dict(slim_run(v), overrides=v['overrides']) for k, v in res['camps'].items()},
    levers={k: [dict(slim_run(d), value=d['value']) for d in lst] for k, lst in res['levers'].items()},
    grid=res['grid_barrier_cc'],
    future={k: dict(slim_run(v), extra_obs=v['extra_obs'], alternatives=[dict(overrides=a['overrides'], alone=r3(a['schools']['site']['p_no_other_ever_in_galaxy']), stuck=r3(a['schools']['site']['posterior']['stuck']),
                                                                             nobody=r3(a['schools']['site']['posterior']['empty_universe'] + a['schools']['site']['posterior']['rare']),
                                                                             unnoticed=r3(a['schools']['site']['posterior']['unnoticed'])) for a in v['alternatives']])
            for k, v in res['future'].items()},
    mc_sd=res['mc_sd'], n_total=res['n_per_run'] * res.get('n_seeds', 8),
)
if os.path.exists(M + '/integrate4_scan.json'):
    sc = json.load(open(M + '/integrate4_scan.json'))
    nb = lambda p: p['empty_universe'] + p['rare']
    out['scan'] = [dict(pa=round(x['prior_alone'], 3), pn=round(nb(x['prior']), 3), site=[round(x['site']['posterior'][c], 3) for c in CLS],
                        a=[round(x[k]['alone'], 3) for k in ('site', 'ssa_all', 'ssa_obs')], nb=[round(nb(x[k]['posterior']), 3) for k in ('site', 'ssa_all', 'ssa_obs')],
                        st=[round(x[k]['posterior']['stuck'], 3) for k in ('site', 'ssa_all', 'ssa_obs')],
                        wf=round(x['params']['w_fast'], 3), qb=round(x['params']['q_barrier'], 3), qc=round(x['params']['q_cc'], 3)) for x in sc]
if os.path.exists(M + '/cells4.json'):
    out['cells'] = json.load(open(M + '/cells4.json'))
# ---- leverage of each measured input (paired runs, models/leverage4.py) ----
if os.path.exists(M + '/leverage4.json'):
    lv = json.load(open(M + '/leverage4.json'))['summary']
    out['leverage'] = {k: dict(a=round(v['site_alone_delta'], 4), s=round(v['site_stuck_delta'], 4), t=round(v['ssa_all_alone_delta'], 4)) for k, v in lv.items()}
# ---- study cards for waves 2 and 3 (hand-written summaries of the refereed results live in report/study_cards.json) ----
if os.path.exists(ROOT + '/report/study_cards.json'):
    out['wave2'] = json.load(open(ROOT + '/report/study_cards.json'))
os.makedirs(ROOT + '/report/site', exist_ok=True)
# ---- working papers (built by tools/build_docs.py) and the reference list (cite + verified link) ----
if os.path.exists(ROOT + '/report/docs_index.json'):
    di = json.load(open(ROOT + '/report/docs_index.json'))
    out['docs'] = di['docs']
    have = {s['key'] for s in di['studies']}
    for c in out.get('wave2', []):
        if c['key'] in have: c['href'] = 'docs/' + c['key'] + '.html'
    out['papers'] = sorted(have)
if os.path.exists(ROOT + '/report/references.json'):
    out['refs'] = [dict(cite=r['cite'].replace('Kipping, D. (2025). Strong', 'Kipping, D. (2025a). Strong').replace('Kipping, D. (2025). Solar', 'Kipping, D. (2025b). Solar'), url=r.get('url') or None)
                   for r in json.load(open(ROOT + '/report/references.json'))]

js = 'window.FERMI = ' + json.dumps(out, separators=(',', ':')) + ';'
open(ROOT + '/report/site/data.js', 'w').write(js)
print('wrote data.js', len(js) // 1024, 'KB; keys:', list(out))