"""Exploratory: by how much would a completeness correction measured with the flare profile
Vasilyev et al. (2024) injected understate the rate of superflares shaped as this study assumes?

Their search smooths each Kepler long-cadence light curve with a 15-point moving average, finds
flares as two consecutive points more than five standard deviations above it, and measures a
flare's energy as the integral of its excess from the first point at five standard deviations, by
linear interpolation, to where the excess falls back to one, divided by 0.75. They corrected their
counts by the fraction of injected flares their search detected; the injected flares had the
average profile of the flares they had found, which puts nearly all its light into two equal
consecutive points. Their frequency figure shows their full sample's flares by energy before and
after that correction, so the ratio of the two is the completeness they applied in each energy bin
(code/read_figures.py reads it, the profile, and their detectability figure from their arXiv
source).

Here a light curve's sensitivity S is the energy whose excess in one 30-minute point equals its
noise, so a flare of energy E puts E/S times its share of light into each point, in units of the
noise. A log-normal distribution of S over light curves and a fraction of flares lost whatever their
size are fitted so that their search, run on their profile injected with white noise less its
moving average, detects it as often as the completeness they applied. The moving average leaves the
flare out, as in their example light curve, where it does not rise through the flare. Flares drawn
with energies even in logarithm, weighted by a power law, either with their profile or with the
Davenport et al. (2014) shape and the widths this study's injections assume (half and two and a half
times them too), at a random phase within a cadence, are searched and measured as they searched and
measured, and each found flare is weighted by one over the completeness they applied at its measured
energy. In each of their two energy bins, the true rate over that estimate is the factor by which a
rate corrected that way falls short for flares of that shape. Their correction divided each
temperature bin's counts by that bin's detected fraction, and their figure's ratio is its net effect
on the whole sample, weighted by where the flares fell; reading that ratio as the response of one
distribution of sensitivities, and taking every light curve to produce flares at the same rate, are
assumptions of this emulation, not of their correction.

It also sets their detectability figure beside the completeness they applied, and integrates their
corrected frequency over their Table 1's two energy bins.

python cadence_test.py <data dir> <results dir>
Writes cadence.json.
"""
import json
import os
import sys

import numpy as np
import pandas as pd
from scipy.ndimage import uniform_filter1d
from scipy.optimize import minimize

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

CADENCE_D = 1765.46 / sf.DAY_S   # Kepler long cadence, 29.4 minutes
SMOOTH = 15                      # their moving average, 7.5 hours
DETECT_Z, END_Z = 5.0, 1.0       # two consecutive points above the first; the energy ends at the second
ENERGY_FRACTION = 0.75           # they divide the measured energy by this
POINTS, PEAK = 61, 20            # each simulated stretch, the flare's peak or first point in cadence PEAK
SEARCH = 4                       # a detection's first point within this many cadences of PEAK
BINS = ((34.0, 35.0), (35.0, np.inf))   # their Table 1's energy bins, log10 erg
LOG_E_RANGE = (32.5, 36.5)       # true energies drawn evenly in log, weighted by the power law
WIDTHS = {'short': 0.5, 'base': 1.0, 'long': 2.5}   # times this study's assumed widths
STEEPER = 0.3                    # a steeper power law, by this much in its cumulative slope
FIT_DRAWS, DRAWS, CHUNK = 40_000, 600_000, 20_000
SEED = 20261025


def noise(rng, n):
    """White noise less its moving average, scaled so that its standard deviation, the σ their
    thresholds are in units of, is one."""
    white = rng.normal(size=(n, POINTS)) * np.sqrt(SMOOTH / (SMOOTH - 1))
    return white - uniform_filter1d(white, SMOOTH, axis=1, mode='nearest')


def davenport_cum(x):
    """The Davenport template's integral from its start to x, in FWHM from the peak."""
    def rise(u):
        return u + 1.941 * u ** 2 / 2 - 0.175 * u ** 3 / 3 - 2.246 * u ** 4 / 4 - 1.125 * u ** 5 / 5
    d = np.maximum(x, 0.0)
    return (rise(np.clip(x, -1.0, 0.0)) - rise(-1.0)
            + 0.6890 / 1.600 * -np.expm1(-1.600 * d) + 0.3030 / 0.2783 * -np.expm1(-0.2783 * d))


def davenport_shares(fwhm_d, phase):
    """Each cadence's share of a Davenport flare's light; the peak falls `phase` of the way into
    cadence PEAK."""
    edges = (np.arange(POINTS + 1)[None, :] - PEAK - phase[:, None]) * CADENCE_D / fwhm_d[:, None]
    return np.diff(davenport_cum(edges), axis=1) / sf.DAVENPORT_AREA


def profile_shares(profile, n):
    """Their profile's share of the light in each cadence, its first point in cadence PEAK; its
    energy is its sum, each point standing for a whole cadence."""
    out = np.zeros((n, POINTS))
    out[:, PEAK + profile.cadence.values] = profile.flux.values / profile.flux.sum()
    return out


def search(excess):
    """Their search of each stretch, in units of σ: whether it finds a flare starting near PEAK,
    and the energy it measures there in units of the sensitivity."""
    hi = excess > DETECT_Z
    pair = (hi[:, :-1] & hi[:, 1:])[:, PEAK - SEARCH:PEAK + SEARCH + 1]
    found = pair.any(axis=1)
    start = PEAK - SEARCH + pair.argmax(axis=1)
    rows = np.arange(len(excess))
    below = (np.arange(POINTS)[None, :] > start[:, None]) & (excess < END_Z)
    end = np.where(below.any(axis=1), below.argmax(axis=1), POINTS - 1)   # first point below 1σ
    last = end - 1
    steps = np.concatenate([np.zeros((len(excess), 1)), np.cumsum((excess[:, 1:] + excess[:, :-1]) / 2, axis=1)], axis=1)
    area = steps[rows, last] - steps[rows, start]
    a, b = excess[rows, last], excess[rows, end]
    tail = np.where(b < END_Z, (a - END_Z) / np.maximum(a - b, 1e-12), 0.0)
    area += (a + END_Z) / 2 * tail
    return found, area / ENERGY_FRACTION


def applied_completeness(frequency):
    """Their full sample's frequency before over after their correction, in each energy bin with
    both, at the bin's center."""
    f = frequency.dropna()
    return pd.DataFrame({'log_e': (f.log_e_lo + f.log_e_hi) / 2, 'r': 10 ** (f.observed - f.corrected)})


def fit_sensitivity(applied, shares, rng):
    """Median log10 sensitivity, its spread in dex, and the fraction lost, that make their search
    detect their profile as often as the completeness they applied."""
    z = rng.normal(size=FIT_DRAWS)
    base = noise(rng, FIT_DRAWS)[:, PEAK - SEARCH:PEAK + SEARCH + 2]
    s = shares[:1, PEAK - SEARCH:PEAK + SEARCH + 2]

    def detected(par):
        mu, log_spread, logit_loss = par
        out = []
        for le in applied.log_e:
            x = 10 ** (le - mu - np.exp(log_spread) * z)
            hi = s * x[:, None] + base > DETECT_Z
            out.append((hi[:, :-1] & hi[:, 1:]).any(axis=1).mean())
        return np.array(out) / (1 + np.exp(logit_loss))

    def cost(par):
        return float(np.sum((detected(par) - applied.r.values) ** 2))

    starts = [(mu, np.log(0.4), -3.0) for mu in np.arange(32.0, 34.1, 0.5)]
    best = min((minimize(cost, x0, method='Nelder-Mead', options={'xatol': 1e-4, 'fatol': 1e-9, 'maxiter': 3000})
                for x0 in starts), key=lambda o: o.fun)
    mu, log_spread, logit_loss = best.x
    fit = detected(best.x)
    return {'bins': int(len(applied)), 'median_log_sensitivity': float(mu), 'spread_dex': float(np.exp(log_spread)),
            'lost': float(1 - 1 / (1 + np.exp(logit_loss))),
            'rms': float(np.sqrt(np.mean((fit - applied.r.values) ** 2))),
            'max_difference': float(np.max(np.abs(fit - applied.r.values)))}


def simulate(fit, profile):
    """For every shape, from the same draws: each flare's true energy, whether their search finds
    it, and the energy it measures (log10)."""
    names = ['two_point'] + list(WIDTHS)
    log_e, out = [], {k: ([], []) for k in names}
    for chunk in range(DRAWS // CHUNK):
        rng = np.random.default_rng([SEED, chunk])
        le = rng.uniform(*LOG_E_RANGE, CHUNK)
        x = 10 ** (le - fit['median_log_sensitivity'] - fit['spread_dex'] * rng.normal(size=CHUNK))
        kept = rng.uniform(size=CHUNK) >= fit['lost']
        phase = rng.uniform(size=CHUNK)
        scatter = 10 ** (ij.SCATTER_DEX * rng.normal(size=CHUNK))
        e = noise(rng, CHUNK)
        log_e.append(le)
        for name in names:
            if name == 'two_point':
                shares = profile_shares(profile, CHUNK)
            else:
                fwhm = ij.T_HALF_34_D * WIDTHS[name] * 10 ** (ij.DURATION_SLOPE * (le - 34.0)) * scatter
                shares = davenport_shares(fwhm, phase)
            found, measured = search(shares * x[:, None] + e)
            found &= kept & (measured > 0)
            out[name][0].append(found)
            out[name][1].append(np.where(found, le - np.log10(x) + np.log10(np.where(found, measured, 1.0)), np.nan))
    return np.concatenate(log_e), {k: (np.concatenate(f), np.concatenate(m)) for k, (f, m) in out.items()}


def factors(log_e, found, log_meas, applied, beta):
    """In each energy bin, the true rate over their estimate of it, for flares of one shape whose
    energies follow a cumulative power law of slope beta."""
    weight = 10 ** (-beta * (log_e - 34.0))
    out = []
    for lo, hi in BINS:
        truth = weight[(log_e >= lo) & (log_e < hi)].sum()
        sel = found & (log_meas >= lo) & (log_meas < hi)
        estimate = np.sum(weight[sel] / np.interp(log_meas[sel], applied.log_e, applied.r))
        out.append(float(truth / estimate))
    return out


def corrected_rates(frequency, beta):
    """Their corrected frequency integrated over their Table 1's two energy bins, per star per year;
    a bin that straddles an edge is shared by the power law."""
    f = frequency.dropna()
    cum = lambda x: 10 ** (-beta * (x - 34.0))  # noqa: E731
    out = []
    for lo, hi in BINS:
        total = 0.0
        for a, b, g in zip(f.log_e_lo, f.log_e_hi, f.corrected):
            u, v = max(a, lo), min(b, hi)
            if v > u:
                total += 10 ** g * (10 ** b - 10 ** a) * (cum(u) - cum(v)) / (cum(a) - cum(b))
        out.append(total)
    return out


def main():
    data, out = sys.argv[1:3]
    meta = os.path.join(data, 'meta')
    detectability = pd.read_csv(os.path.join(meta, 'vasilyev_detectability.csv'))
    profile = pd.read_csv(os.path.join(meta, 'vasilyev_profile.csv'))
    frequency = pd.read_csv(os.path.join(meta, 'vasilyev_frequency.csv'))
    rates = an.VASILYEV['all']
    beta = an.power_law(**rates)[1]
    applied = applied_completeness(frequency)

    fit = fit_sensitivity(applied, profile_shares(profile, 1), np.random.default_rng(SEED))
    log_e, runs = simulate(fit, profile)
    res = {'calibration': fit, 'beta': beta, 'beta_steeper': beta + STEEPER}
    for name, (found, log_meas) in runs.items():
        low, high = factors(log_e, found, log_meas, applied, beta)
        res[name] = {'factor_low_bin': low, 'factor_high_bin': high}
    low, high = factors(log_e, *runs['base'], applied, beta + STEEPER)
    res['base_steeper'] = {'factor_low_bin': low, 'factor_high_bin': high}

    # Their detectability figure beside the completeness they applied: the lowest energy from which
    # the applied completeness exceeds the figure's at every temperature
    figure = detectability.pivot(index='log_e', columns='teff', values='r')
    best = np.array([max(np.interp(le, figure.index, figure[t]) for t in figure.columns) for le in applied.log_e])
    above = applied.r.values > best
    first = len(above) - np.argmin(above[::-1]) if not above.all() else 0
    near = int(np.argmin(np.abs(applied.log_e.values - 34.95)))
    res['figure'] = {'applied_exceeds_from_log_e': float(applied.log_e.iloc[first]) if first < len(above) else None,
                     'log_e_near_1e35': float(applied.log_e.iloc[near]), 'applied_r_near_1e35': float(applied.r.iloc[near]),
                     'figure_max_r_near_1e35': float(best[near])}
    low, high = corrected_rates(frequency, beta)
    res['rates'] = {'figure_low_bin': low, 'figure_high_bin': high,
                    'table_low_bin': rates['f34_35'], 'table_high_bin': rates['f35']}
    with open(os.path.join(out, 'cadence.json'), 'w') as fh:
        json.dump(res, fh, indent=1, sort_keys=True)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()