#!/usr/bin/env python3
"""Frozen observed OGLE I-band magnitudes; weighted floating-mean full-grid search.
Python >=3.10 and numpy. The source values remain measured, never replaced by the fit.
"""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np

SCRIPT_DIRECTORY = Path(__file__).resolve().parent
ROOT = SCRIPT_DIRECTORY.parents[2] if len(SCRIPT_DIRECTORY.parents) >= 3 else SCRIPT_DIRECTORY
REPOSITORY_SOURCE = ROOT/'src/features/labs/observing-cadence/data/ogle-lmc-rrlyr-00001.dat'
SOURCE = REPOSITORY_SOURCE if REPOSITORY_SOURCE.exists() else SCRIPT_DIRECTORY/'ogle-lmc-rrlyr-00001.dat'
SOURCE_HASH = '6d6bb05178881e421ca2c03cf77bc100d51f74c87efbd523a9f017216c6e5a1b'

def projection_basis(time, sigma, frequencies):
    constant = 1/sigma
    constant /= np.linalg.norm(constant)
    phase = 2*np.pi*frequencies[:,None]*time[None,:]
    cosine, sine = np.cos(phase)/sigma, np.sin(phase)/sigma
    cosine -= (np.einsum("ij,j->i",cosine,constant,optimize=False))[:,None]*constant
    sine -= (np.einsum("ij,j->i",sine,constant,optimize=False))[:,None]*constant
    cosine_norm = np.linalg.norm(cosine,axis=1)
    good_cosine = cosine_norm > 1e-10
    cosine[good_cosine] /= cosine_norm[good_cosine,None]
    cosine[~good_cosine] = 0
    sine -= np.sum(sine*cosine,axis=1)[:,None]*cosine
    sine_norm = np.linalg.norm(sine,axis=1)
    good_sine = sine_norm > 1e-10
    sine[good_sine] /= sine_norm[good_sine,None]
    sine[~good_sine] = 0
    return constant,cosine,sine

def powers(basis, whitened_values):
    constant,cosine,sine = basis
    centered = whitened_values-constant[:,None]*(np.einsum("i,ij->j",constant,whitened_values,optimize=False))[None,:]
    reference = np.sum(centered**2,axis=0)
    return ((np.einsum("ij,jk->ik",cosine,centered,optimize=False))**2+(np.einsum("ij,jk->ik",sine,centered,optimize=False))**2)/reference

def analyze(data,frequencies,seed):
    time,value,sigma = data.T
    basis = projection_basis(time,sigma,frequencies)
    measured_power = powers(basis,(value/sigma)[:,None])[:,0]
    best_index = int(np.argmax(measured_power))
    random = np.random.default_rng(seed)
    null_trials = 4096
    null_maxima = powers(basis,random.normal(size=(len(time),null_trials))).max(axis=0)
    threshold = float(np.quantile(null_maxima,.95,method='higher'))
    heldout_trials = 1024
    heldout_maxima = powers(basis,random.normal(size=(len(time),heldout_trials))).max(axis=0)
    exceedances = int(np.count_nonzero(null_maxima>=measured_power[best_index]))
    # Permute standardized observed deviations, a separate exchangeability test.
    standardized = (value-np.average(value,weights=1/sigma**2))/sigma
    permutations = np.stack([random.permutation(standardized) for _ in range(256)],axis=1)
    permutation_maxima = powers(basis,permutations).max(axis=0)
    injected_frequencies = random.uniform(.3,4.8,128)
    phases = random.uniform(0,2*np.pi,128)
    injection = .1*np.sin(2*np.pi*time[:,None]*injected_frequencies[None,:]+phases[None,:])/sigma[:,None]+random.normal(size=(len(time),128))
    injection_power = powers(basis,injection)
    recovered_indices = np.argmax(injection_power,axis=0)
    recovered_frequencies = frequencies[recovered_indices]
    detected = injection_power[recovered_indices,np.arange(128)]>=threshold
    tolerance = 1/(time.max()-time.min())
    correct = np.abs(recovered_frequencies-injected_frequencies)<tolerance
    return {
        'samples': len(time),
        'span_days': float(np.ptp(time)),
        'frequency_grid_per_day': frequencies.tolist(),
        'power': measured_power.tolist(),
        'best_frequency_per_day': float(frequencies[best_index]),
        'best_period_days': float(1 / frequencies[best_index]),
        'peak_power': float(measured_power[best_index]),
        'window_power': np.abs(np.exp(2j * np.pi * frequencies[:, None] * time[None, :]).mean(axis=1)) ** 2,
        'null_calibration': {
            'kind': 'constant-source independent Gaussian errors using quoted magnitude sigmas',
            'seed': seed,
            'trials': null_trials,
            'maxima': null_maxima.tolist(),
            'threshold_5percent': threshold,
            'exceedances': exceedances,
            'finite_mc_tail_rank': (1 + exceedances) / (null_trials + 1),
            'zero_exceedance_95percent_upper': float(1 - 0.05 ** (1 / null_trials)) if exceedances == 0 else None,
            'heldout_trials': heldout_trials,
            'heldout_exceedances': int(np.count_nonzero(heldout_maxima >= threshold)),
            'heldout_rate': float(np.mean(heldout_maxima >= threshold)),
        },
        'permutation_test': {
            'kind': 'standardized observed deviations permuted at fixed times; exchangeability assumed',
            'trials': 256,
            'finite_mc_tail_rank': float((1 + np.count_nonzero(permutation_maxima >= measured_power[best_index])) / 257),
        },
        'injection_recovery': {
            'kind': 'synthetic sinusoidal magnitude injections on measured times/errors, separate from observed magnitudes',
            'trials': 128,
            'semi_amplitude_mag': 0.1,
            'frequency_range_per_day': [0.3, 4.8],
            'correct_frequency_tolerance_per_day': float(tolerance),
            'detected': int(detected.sum()),
            'correct_and_detected': int((correct & detected).sum()),
            'alias_and_detected': int((~correct & detected).sum()),
            'missed': int((~detected).sum()),
            'true_frequencies': injected_frequencies.tolist(),
            'recovered_frequencies': recovered_frequencies.tolist(),
        },
        'observations': [{'time_from_first_days': float(t), 'magnitude': float(v), 'sigma_mag': float(s)} for (t, v, s) in data],
    }

def clean(value):
    if isinstance(value,np.ndarray):return value.tolist()
    if isinstance(value,dict):return {key:clean(item) for key,item in value.items()}
    if isinstance(value,list):return [clean(item) for item in value]
    return value

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=Path, default=SOURCE)
    parser.add_argument('--out', type=Path, default=ROOT/'docs/astronomy/ac07/ogle-replay.json' if REPOSITORY_SOURCE.exists() else SCRIPT_DIRECTORY/'ogle-replay-rerun.json')
    args = parser.parse_args()
    SOURCE = args.input
    if hashlib.sha256(SOURCE.read_bytes()).hexdigest()!=SOURCE_HASH:raise ValueError('Frozen source hash changed')
    source=np.loadtxt(SOURCE)
    selected=source[(source[:,0]<=source[0,0]+160)&np.isfinite(source).all(axis=1)&(source[:,2]>0)].copy()
    origin=float(selected[0,0]);selected[:,0]-=origin
    step=1/(5*np.ptp(selected[:,0]));frequencies=.2+step*np.arange(int((6-.2)/step)+1)
    results={
        'object': 'OGLE-LMC-RRLYR-00001',
        'kind': 'measured OGLE-IV I-band magnitude analysis',
        'input_sha256': SOURCE_HASH,
        'time_convention': 'HJD minus 2450000 in source; subtract first selected value for fitting. Not BJD/TDB.',
        'origin_hjd_minus_2450000': origin,
        'selection': 'all finite positive-sigma rows within 160 days of first point; retain original quoted uncertainties; no clipping',
        'catalog_period_days': 0.6347521,
        'numpy_version': np.__version__,
        'grid_rule': '0.2 to at most 6 cycles/day, step 1/(5*selected baseline); exactly the same grid reused in every null, permutation and injection search',
        'full': analyze(selected, frequencies, 707),
        'every_third': analyze(selected[::3], frequencies, 708),
        'scope': 'Single-harmonic fit in magnitudes. RRab shape is not sinusoidal. Noise-only calibration is conditional on independent Gaussian quoted errors; no red-noise/crowding/systematic model. FAP does not give correct-period probability. Sparse data receive their own null calibration at their own times with the same declared search grid.',
    }
    output=args.out;output.parent.mkdir(parents=True,exist_ok=True);output.write_text(json.dumps(clean(results),indent=2,allow_nan=False)+'\n')
    print(json.dumps({key:{'samples':results[key]['samples'],'best_period_days':results[key]['best_period_days'],'peak_power':results[key]['peak_power'],'null_fap_rank':results[key]['null_calibration']['finite_mc_tail_rank'],'null_5percent_threshold':results[key]['null_calibration']['threshold_5percent'],'heldout_false_alarm_rate':results[key]['null_calibration']['heldout_rate'],'permutation_rank':results[key]['permutation_test']['finite_mc_tail_rank'],'injection':{k:v for k,v in results[key]['injection_recovery'].items() if k not in ['true_frequencies','recovered_frequencies']}} for key in ['full','every_third']},indent=2))