"""From the search's tables to the study's numbers: how many superflares TESS saw on Sun-like stars,
how many the Kepler rate of Vasilyev et al. (2024) predicts it would see, and their ratio.

python analyze.py <data dir> <run dir> <results dir> [--set main|dev] [--localization FILE]
"""
import argparse
import json
import os
import sys

import numpy as np
import pandas as pd
from scipy import stats

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sflare as sf  # noqa: E402


def table(run, name):
    p = os.path.join(run, name + '.csv')
    return pd.read_csv(p if os.path.exists(p) else p + '.gz')

YEAR_D = 365.25
GRID_STEP = 0.2

# Vasilyev et al. (2024), Table 1: superflare frequency per star per year
VASILYEV = {
    'all': {'f34_35': 8.63e-3, 'f35': 3.90e-4},       # 5000-6500 K, Prot > 20 d or unknown
    'sunlike': {'f34_35': 9.31e-3, 'f35': 3.59e-4},   # 5500-6000 K, Rvar < 0.3 per cent
}
# The earlier rate Vasilyev et al. quote for Sun-like stars: superflares above 1e34 erg once
# per 2000-3000 years (Notsu et al. 2019; Okamoto et al. 2021); 2500 years, same shape.
EARLIER_F34 = 1 / 2500.0
# Star-level rules, fixed in the plan
LS_POWER, LS_AMP = 0.3, 1e-3       # TESS fast rotator: Lomb-Scargle peak 0.2-12 d above both
ECLIPSE_ED_S, ECLIPSE_N = 20.0, 2  # eclipser: at least 2 dips of ED 20 s or more in a sector
SUNLIKE_TEFF = (5500, 6000)
SUNLIKE_RVAR = 3e-3
COINCIDENCE_WINDOW_D = 3 * sf.CADENCE_S / sf.DAY_S
COINCIDENCE_MAX = 1                # a counted flare has at most this many other stars coincident


def power_law(f34_35, f35):
    """Cumulative F(>E) = F34 (E/1e34)^-beta through the two numbers of a Table 1 row."""
    f34 = f34_35 + f35
    return f34, float(np.log10(f34 / f35))


def bin_mass(log_e, f34, beta):
    """The model's flares per star-year in the 0.2-dex bin centred on log_e; the top grid bin
    (35.6) takes the whole tail above its lower edge."""
    log_e = np.asarray(log_e, float)
    lo = log_e - GRID_STEP / 2
    hi = log_e + GRID_STEP / 2
    cum = lambda x: f34 * 10 ** (-beta * (x - 34.0))  # noqa: E731
    return np.where(log_e >= 35.6 - 1e-9, cum(lo), cum(lo) - cum(hi))


def poisson_ci(n, cl=0.95):
    a = 1 - cl
    lo = 0.0 if n == 0 else stats.chi2.ppf(a / 2, 2 * n) / 2
    hi = stats.chi2.ppf(1 - a / 2, 2 * n + 2) / 2
    return float(lo), float(hi)


def coincidences(c):
    """For each candidate: other stars with a candidate of either sign in the same sector and
    camera within three cadences of its peak."""
    c = c.reset_index(drop=True)
    n_other = np.zeros(len(c), int)
    for _, g in c.groupby(['sector', 'camera']):
        g = g.sort_values('tpk')
        t, tic, idx = g.tpk.values, g.tic.values, g.index.values
        lo = np.searchsorted(t, t - COINCIDENCE_WINDOW_D, 'left')
        hi = np.searchsorted(t, t + COINCIDENCE_WINDOW_D, 'right')
        for k in range(len(g)):
            n_other[idx[k]] = len(set(tic[lo[k]:hi[k]]) - {tic[k]})
    c['n_coincident'] = n_other
    return c


class Study:
    def __init__(self, data, run, which, localization=None):
        stars = pd.read_csv(os.path.join(data, 'sample_stars.csv'))
        self.stars = stars[stars.dev == (which == 'dev')].copy()
        lcs = table(run, 'lcs')
        lcs = lcs[lcs.tic.isin(set(self.stars.tic))].copy()
        cands = table(run, 'cands')
        cands = coincidences(cands[cands.file.isin(set(lcs.file))])
        inj = table(run, 'inj')
        if 'model' not in inj:
            inj['model'] = 'base'
        self.resp = os.path.join(data, 'tess-response-function.csv')

        fast = set(lcs.tic[(lcs.ls_power > LS_POWER) & (lcs.ls_amp > LS_AMP)])
        deep = cands[(cands.kind == 'dip') & (cands.ed >= ECLIPSE_ED_S)].groupby('file').size()
        eclipser = set(lcs.tic[lcs.file.isin(set(deep[deep >= ECLIPSE_N].index))])
        s = self.stars
        s['tess_fast'] = s.tic.isin(fast)
        s['eclipser'] = s.tic.isin(eclipser)
        s['rvar'] = s.tic.map(lcs.groupby('tic').rvar.median())
        s['excluded'] = s.tess_fast | s.eclipser
        self.keep = set(s.tic[~s.excluded])
        self.use = lcs[lcs.tic.isin(self.keep)].copy()
        self.use['years'] = self.use.days_good / YEAR_D
        self.years = dict(zip(self.use.file, self.use.years))
        self.inj = inj[inj.file.isin(set(self.use.file))]
        self.grid = np.sort(self.inj.log_e.unique())
        if localization and os.path.exists(localization):
            loc = pd.read_csv(localization)
            cands = cands.merge(loc[['file', 'tpk', 'on_target']], on=['file', 'tpk'], how='left')
        else:
            cands['on_target'] = np.nan
        self.cands = cands[cands.tic.isin(self.keep)]

    # ----------------------------------------------------------------------------------------
    def counted(self, kind, files, e_star=1e34, scale=None, localized=False):
        d = self.cands[(self.cands.kind == kind) & self.cands.file.isin(files)]
        energy = d.energy * (d.tic.map(scale) if scale is not None else 1.0)
        ok = d.vetted.astype(bool) & (energy >= e_star) & (d.n_coincident <= COINCIDENCE_MAX)
        if localized:
            ok &= d.on_target.fillna(True).astype(bool)
        return d[ok].assign(energy_used=energy[ok])

    def expected(self, files, f34, beta, thr='340', model='base', shift=None):
        """Model flares per star-year in each injection bin, times the years of the light curve,
        times the fraction of injections there recovered with measured energy over threshold.
        thr may be a per-file mapping (column names), shift a per-file log10 energy scale."""
        sub = self.inj[(self.inj.model == model) & self.inj.file.isin(files)]
        if isinstance(thr, dict):
            names = sorted({thr[f] for f in sub.file})
            table = sub[[f'n_rec_{n}' for n in names]].values.astype(float)
            col = sub.file.map({f: names.index(thr[f]) for f in set(sub.file)}).values
            rec = table[np.arange(len(sub)), col]
        else:
            rec = sub[f'n_rec_{thr}'].values.astype(float)
        frac = np.where(sub.n_inj.values > 0, rec / np.maximum(sub.n_inj.values, 1), 0.0)
        log_true = sub.log_e.values + (sub.file.map(shift).values if shift is not None else 0.0)
        mass = bin_mass(np.round(sub.log_e.values, 1), f34, beta) * 10 ** (-beta * (log_true - sub.log_e.values))
        return float(np.sum(frac * sub.file.map(self.years).values * mass))

    def test(self, tics, label, ref='all', e_star=1e34, thr='340', localized=False):
        files = set(self.use.file[self.use.tic.isin(tics)])
        f34, beta = power_law(**VASILYEV[ref])
        n_v = self.expected(files, f34, beta, thr)
        n_e = self.expected(files, EARLIER_F34, beta, thr)
        obs = self.counted('flare', files, e_star, localized=localized)
        n = len(obs)
        lo, hi = poisson_ci(n)
        return {
            'label': label, 'stars': int(self.use.tic[self.use.file.isin(files)].nunique()), 'light_curves': len(files),
            'star_years': float(self.use.years[self.use.file.isin(files)].sum()), 'observed': n,
            'expected_vasilyev': n_v, 'expected_earlier': n_e,
            'ratio': n / n_v if n_v > 0 else None, 'ratio_ci95': [lo / n_v, hi / n_v] if n_v > 0 else None,
            'p_at_most_observed_if_vasilyev': float(stats.poisson.cdf(n, n_v)),
            'p_at_least_observed_if_earlier': float(stats.poisson.sf(n - 1, n_e)) if n > 0 else 1.0,
            'rate_above_1e34_per_year': (n / n_v) * f34 if n_v > 0 else None,
            'rate_above_1e34_ci95': [lo / n_v * f34, hi / n_v * f34] if n_v > 0 else None,
            'years_per_superflare': (n_v / n) / f34 if n > 0 else None,
            'vasilyev_f34': f34, 'vasilyev_beta': beta,
            'flares': obs[['file', 'tic', 'sector', 'tpk', 'energy', 'ed', 'fwhm', 'z_pk', 'n_coincident']].to_dict('records'),
        }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('data'); ap.add_argument('run'); ap.add_argument('results')
    ap.add_argument('--set', default='main')
    ap.add_argument('--localization')
    a = ap.parse_args()
    os.makedirs(a.results, exist_ok=True)
    st = Study(a.data, a.run, a.set, a.localization)
    s, keep, use = st.stars, st.keep, st.use
    out = {}
    out['primary'] = st.test(keep, 'Sun-like stars, 5000-6500 K')

    # 1. Sun-like row
    sun = set(s.tic[(s.Teff > SUNLIKE_TEFF[0]) & (s.Teff < SUNLIKE_TEFF[1]) & (s.rvar < SUNLIKE_RVAR) & ~s.excluded])
    out['sunlike'] = st.test(sun, 'Teff 5500-6000 K, Rvar below 0.3 per cent', ref='sunlike')
    # 2. Activity
    act = s.set_index('tic').activityindex_espcs
    med = float(np.nanmedian(act[list(keep)]))
    out['low_activity'] = st.test({t for t in keep if act[t] < med}, f'Gaia activity index below the median, {med:.4f} nm')
    out['high_activity'] = st.test({t for t in keep if act[t] >= med}, 'Gaia activity index at or above the median')
    # 3. Brightness
    tm = s.set_index('tic').tmag_list
    for lo_, hi_ in ((0, 8.5), (8.5, 9.5), (9.5, 10.5)):
        out[f'tmag_{lo_}_{hi_}'] = st.test({t for t in keep if lo_ <= tm[t] < hi_}, f'Tmag {lo_}-{hi_}')
    # 4. Other thresholds
    out['above_10^34.5'] = st.test(keep, 'measured energies of 10^34.5 erg or more', e_star=10 ** 34.5, thr='345')
    out['above_10^35'] = st.test(keep, 'measured energies of 10^35 erg or more', e_star=1e35, thr='350')
    # 5. Background from the dips
    files = set(use.file)
    p = out['primary']
    dips = st.counted('dip', files)
    excess = p['observed'] - len(dips)
    sd = np.sqrt(p['observed'] + len(dips))
    out['background_subtracted'] = {
        'dips': int(len(dips)), 'excess': int(excess),
        'ratio': excess / p['expected_vasilyev'], 'ratio_approx_ci95': [(excess - 1.96 * sd) / p['expected_vasilyev'],
                                                                       (excess + 1.96 * sd) / p['expected_vasilyev']],
        'dip_events': dips[['file', 'tpk', 'energy']].to_dict('records')}
    # 6. Flare durations: the light curves injected with all three, scaled to the whole set
    f34, beta = power_law(**VASILYEV['all'])
    sens = set(st.inj.file[st.inj.model == 'long'])
    if sens:
        base_sub = st.expected(sens, f34, beta, model='base')
        out['durations'] = {'light_curves': len(sens)}
        for m in ('short', 'long'):
            n_v = p['expected_vasilyev'] * st.expected(sens, f34, beta, model=m) / base_sub
            lo, hi = poisson_ci(p['observed'])
            out['durations'][m] = {'expected_vasilyev': n_v, 'ratio': p['observed'] / n_v, 'ratio_ci95': [lo / n_v, hi / n_v]}
    # 7. Flare temperature: our energies on another scale, Vasilyev et al.'s numbers unchanged
    out['flare_temperature'] = {}
    teff, rad = s.set_index('tic').Teff, s.set_index('tic').Rad
    for t_fl in (8000.0, 10000.0):
        scale = {t: sf.energy_per_ed(teff[t], rad[t], st.resp, t_fl) / sf.energy_per_ed(teff[t], rad[t], st.resp) for t in keep}
        fshift = {f: float(np.log10(scale[t])) for f, t in zip(use.file, use.tic)}
        thr = {f: str(int(round(10 * (34.0 - fshift[f])))) for f in use.file}
        if not all(f'n_rec_{v}' in st.inj for v in set(thr.values())):
            continue
        n_v = st.expected(files, f34, beta, thr=thr, shift=fshift)
        n = len(st.counted('flare', files, 1e34, scale=scale))
        lo, hi = poisson_ci(n)
        out['flare_temperature'][str(int(t_fl))] = {'observed': n, 'expected_vasilyev': n_v, 'ratio': n / n_v,
                                                    'ratio_ci95': [lo / n_v, hi / n_v]}
    # 8. Localization in the target pixel files
    if a.localization:
        out['localized'] = st.test(keep, 'counting only flares that locate on their targets', localized=True)
    # 9. Strict single stars
    single = set(s.tic[(s.non_single_star.fillna(0) == 0) & ~(s.rv_chisq_pvalue < 0.01) & ~s.excluded])
    out['strict_single'] = st.test(single, 'no Gaia non-single-star solution or radial-velocity variability')

    out['sample'] = {
        'stars_selected': int(len(s)), 'stars_tess_fast_rotators': int(s.tess_fast.sum()),
        'stars_eclipsers': int(s.eclipser.sum()), 'stars_used': int(len(keep)),
        'light_curves_used': int(len(use)), 'star_years': float(use.years.sum()),
        'median_sigma_ppm': float(use.sigma_ppm.median()),
        'flares_vetted_all_energies': int(st.cands[(st.cands.kind == 'flare')].vetted.astype(bool).sum()),
        'activity_median_nm': med,
    }
    for k, v in out.items():
        with open(os.path.join(a.results, f'{k}.json'), 'w') as fh:
            json.dump(v, fh, indent=1, sort_keys=True, default=float)
    print(json.dumps({k: p[k] for k in ('stars', 'light_curves', 'star_years', 'observed', 'expected_vasilyev',
                                        'expected_earlier', 'ratio', 'ratio_ci95')}, indent=1))


if __name__ == '__main__':
    main()