"""Analyses added after the registered results were seen (exploratory): how the counted flares
are spread over the stars, and which of their stars have neighbors in Gaia DR3 close enough to
share their TESS pixels. code/width_check.py compares the flares' widths.

python exploratory.py <data dir> <run dir> <results dir> <counted flares json>
Writes R14 (spread over stars) and R16 (neighbors).
"""
import json
import os
import sys

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

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

data, run, out, counted = sys.argv[1:5]
st = an.Study(data, run, 'main')
fl = pd.DataFrame(json.load(open(counted))['flares'])

# Each star's expected count under the Kepler rate: its light curves' years times the recovery of
# injections at each energy, summed over the rate's energy bins, as Eq. (2) sums them
f34, beta = an.power_law(**an.VASILYEV['all'])
inj = st.inj[st.inj.model == 'base']
frac = np.where(inj.n_inj > 0, inj.n_rec_340 / inj.n_inj.clip(lower=1), 0.0)
contrib = frac * inj.file.map(st.years).values * an.bin_mass(inj.log_e.round(1).values, f34, beta)
per_lc = pd.Series(contrib, index=inj.file.values).groupby(level=0).sum()
mu_k = per_lc.groupby(st.use.set_index('file').tic).sum()
mu_k = mu_k.reindex(sorted(set(st.use.tic))).fillna(0.0)
n = fl.groupby('tic').size().reindex(mu_k.index).fillna(0).astype(int)
total = int(n.sum())

# R14: repeat flaring against one rate shared by every star, passed through the search
mu = mu_k * total / mu_k.sum()
p2 = stats.poisson.sf(1, mu.values)
# The number of stars with two or more is a sum of independent Bernoulli trials: its exact
# distribution, truncated far above anything it reaches
dist = np.zeros(61)
dist[0] = 1.0
for q in p2:
    dist[1:] = dist[1:] * (1 - q) + dist[:-1] * q
    dist[0] *= 1 - q
obs_repeats = int((n >= 2).sum())
r14 = {
    'stars_with_flares': int((n > 0).sum()),
    'stars_with_two_or_more': obs_repeats,
    'expected_two_or_more_if_one_rate': round(float(p2.sum()), 2),
    'log10_p_two_or_more_if_one_rate': round(float(np.log10(dist[obs_repeats:].sum())), 1),
    'flares_on_stars_with_two_or_more': int(n[n >= 2].sum()),
    'most_flares_on_one_star': int(n.max()),
}
# Stars with at least one counted flare, against the number the Kepler rate predicts would have
# one, and the multiple of that rate that would give as many if every star shared it
occ = int((n > 0).sum())
expected_occ = float(np.sum(1 - np.exp(-mu_k.values)))


def occupied(rho):
    return float(np.sum(1 - np.exp(-rho * mu_k.values)))


rho = optimize.brentq(lambda x: occupied(x) - occ, 1e-3, 1e3)


def z(rho_, side):
    p = 1 - np.exp(-rho_ * mu_k.values)
    return (occ - p.sum()) / np.sqrt(np.sum(p * (1 - p))) - side


lo = optimize.brentq(lambda x: z(x, 1.96), 1e-3, rho)
hi = optimize.brentq(lambda x: z(x, -1.96), rho, 1e3)
r14.update({'expected_stars_with_flares_kepler': round(expected_occ, 1),
            'occupancy_rate_multiple': round(rho, 2), 'occupancy_rate_multiple_low': round(lo, 2),
            'occupancy_rate_multiple_high': round(hi, 2)})

# R16: Gaia DR3 neighbors within one TESS pixel, flaring stars against other stars searched
nb = pd.read_csv(os.path.join(data, 'meta', 'gaia_neighbors.csv.gz'))
g = st.stars.set_index('tic').phot_g_mean_mag
nb = nb[nb.tic.isin(g.index)]
nb['dG'] = nb.Gmag - nb.tic.map(g)
ctl = pd.read_csv(os.path.join(data, 'meta', 'neighbor_controls.csv'))
flaring = set(fl.tic)
controls = (set(ctl.tic) - flaring) & st.keep
r16 = {'flaring_stars': len(flaring), 'control_stars': len(controls)}
for name, sel in (('any_dG8', (nb.angDist < 21) & (nb.dG < 8)),
                  ('red_dG8', (nb.angDist < 21) & (nb.dG < 8) & (nb['BP-RP'] > 1.8))):
    has = set(nb.tic[sel])
    a, b = len(flaring & has), len(controls & has)
    r16[name] = {'flaring': a, 'flaring_fraction': round(a / len(flaring), 2),
                 'controls': b, 'control_fraction': round(b / len(controls), 2),
                 'p_fisher': float(f"{stats.fisher_exact([[a, len(flaring) - a], [b, len(controls) - b]])[1]:.2g}")}
for k, v in (('R14', r14), ('R16', r16)):
    with open(os.path.join(out, k + '.json'), 'w') as fh:
        json.dump(v, fh, indent=1, sort_keys=True)
print(json.dumps({'R14': r14, 'R16': r16}, indent=1))