#!/usr/bin/env python3
"""Known-truth activity/reflex/offset experiments on the real 71-point sampling.

All velocities and indicators generated here are SYNTHETIC. The source table
supplies times and quoted noise scales only. There are no new real instruments.
"""
import argparse,json
from pathlib import Path
import numpy as np
from rv_kernel import linear_gaussian_posterior
from rv_posterior import read_observations
from surface_reference import observe


def coefficient_summary(mean,covariance,labels):
    return dict(labels=labels,mean=mean.tolist(),full_covariance=covariance.tolist(),
                signed_sine_interval_95=[float(mean[-2]-1.96*np.sqrt(covariance[-2,-2])),float(mean[-2]+1.96*np.sqrt(covariance[-2,-2]))])


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--data',type=Path,default=Path('.'))
    parser.add_argument('--out',type=Path,default=Path('injection-results.json'));args=parser.parse_args()
    rows,times,observed,errors,normalization=read_observations(args.data)
    count=len(times);random=np.random.default_rng(1815)
    periodic=[]
    for planet_period in [23.81,12.7,8.0]:
        planet_sine=np.sin(2*np.pi*times/planet_period);planet_cosine=np.cos(2*np.pi*times/planet_period)
        active_sine=np.sin(2*np.pi*times/23.81);active_cosine=np.cos(2*np.pi*times/23.81)
        for truth in ['planet_only','activity_only']:
            planet_truth=6. if truth=='planet_only' else 0.
            activity_truth=0. if truth=='planet_only' else 8.
            rv_values=planet_truth*planet_sine+activity_truth*active_sine+random.normal(size=count)*errors[:count]
            calcium_values=.3*activity_truth*active_sine+random.normal(size=count)*errors[count:]
            for activity_prior in [5.,20.,60.]:
                for calcium_used in [False,True]:
                    design=np.zeros((2*count if calcium_used else count,6))
                    design[:count,0]=1;design[:count,2]=planet_sine;design[:count,3]=planet_cosine
                    design[:count,4]=active_sine;design[:count,5]=active_cosine
                    if calcium_used:
                        design[count:,1]=1;design[count:,4]=.3*active_sine;design[count:,5]=.3*active_cosine
                    else:
                        design[:,1]=0
                    values=np.r_[rv_values,calcium_values] if calcium_used else rv_values
                    noise=np.diag(errors**2 if calcium_used else errors[:count]**2)
                    _,mean,covariance=linear_gaussian_posterior(values,noise,design,[100.,10.,10.,10.,activity_prior,activity_prior])
                    periodic.append(dict(kind='single-harmonic GP (random Gaussian sine/cosine), coherent limit; not the full quasi-periodic real-data model',
                        truth=truth,planet_period_days=planet_period,rotation_days=23.81,planet_sine_truth_m_s=planet_truth,
                        activity_sine_truth_m_s=activity_truth,activity_prior_scale_m_s=activity_prior,calcium_used=calcium_used,
                        proxy_scale='.3 normalized-calcium units per metre/second; prescribed synthetic response',
                        planet_sine_mean_m_s=float(mean[2]),planet_sine_standard_deviation_m_s=float(np.sqrt(covariance[2,2])),
                        activity_sine_mean_m_s=float(mean[4]),full_mean=mean.tolist(),full_covariance=covariance.tolist()))
    inputs=json.loads((args.data/'surface-input.json').read_text())
    settings=next(entry['settings'] for entry in inputs['cases'] if entry['name']=='activity')
    physical=[observe(settings,float(time),80,320) for time in times]
    values=np.array([sample['centroid_m_s'] for sample in physical])+random.normal(size=count)*errors[:count]
    frequency_grid=np.linspace(1/35,1/5,301)
    grid=[];best=None
    for frequency in frequency_grid:
        design=np.column_stack([np.ones(count),np.sin(2*np.pi*times*frequency),np.cos(2*np.pi*times*frequency)])
        whitened=design/errors[:count,None];normalized=values/errors[:count]
        coefficients=np.linalg.lstsq(whitened,normalized,rcond=None)[0]
        residual=normalized-np.einsum('ij,j->i',whitened,coefficients)
        objective=float(np.sum(residual**2));grid.append(dict(period_days=float(1/frequency),chi_squared=objective))
        if best is None or objective<best[0]:best=(objective,frequency,design)
    _,frequency,design=best
    _,mean,covariance=linear_gaussian_posterior(values,np.diag(errors[:count]**2),design,[100.,10.,10.])
    pure_activity=dict(kind='physical rotating Gaussian dark region, no companion, independently polar-integrated; all velocities synthetic',
        settings=settings,period_search=grid,selected_period_days=float(1/frequency),
        false_kepler_amplitude_m_s=float(np.hypot(mean[1],mean[2])),
        selected_period_posterior=coefficient_summary(mean,covariance,['offset_m_s','sine_m_s','cosine_m_s']),
        selection_scope='conditional posterior after a frequency search; no calibrated false-alarm or planet probability',
        samples=[dict(time_offset_days=float(time),noise_error_m_s=float(error),synthetic_rv_m_s=float(value),
            true_activity_rv_m_s=sample['centroid_m_s'],true_planet_rv_m_s=0.,projected_active_proxy=sample['projected_active_proxy'])
            for time,error,value,sample in zip(times,errors[:count],values,physical)])
    instrument=(times>=12).astype(float)
    values=12*instrument+random.normal(size=count)*errors[:count]
    sine=np.sin(2*np.pi*times/20);cosine=np.cos(2*np.pi*times/20)
    offset_results=[]
    for separate in [False,True]:
        design=np.column_stack([np.ones(count),instrument,sine,cosine]) if separate else np.column_stack([np.ones(count),sine,cosine])
        scales=[100.,100.,10.,10.] if separate else [100.,10.,10.]
        _,mean,covariance=linear_gaussian_posterior(values,np.diag(errors[:count]**2),design,scales)
        offset_results.append(dict(separate_instrument_offsets=separate,planet_truth_m_s=0.,offset_step_m_s=12.,
            amplitude_m_s=float(np.hypot(mean[-2],mean[-1])),full_mean=mean.tolist(),full_covariance=covariance.tolist()))
    output=dict(schema='ac18-known-truth-injections-v1',kind='synthetic physical surface, single-harmonic latent GP and artificial instrument split; actual HARPS sampling/noise scales only',
        seed=1815,periodic_gp_experiments=periodic,physical_activity_false_positive=pure_activity,
        synthetic_instrument_split=offset_results,
        boundaries=['calcium scaling in the injections is prescribed, not a measured chromospheric calibration',
            'all71realrowsremainHARPS; syntheticinstrument labels are not historical instruments',
            'signed-coefficient credible intervals are conditional on their Gaussian models and priors, not detection probabilities'])
    args.out.write_text(json.dumps(output,indent=2,allow_nan=False)+'\n')
    print(json.dumps(dict(false_kepler_amplitude=pure_activity['false_kepler_amplitude_m_s'],false_period=pure_activity['selected_period_days'],offset_experiments=offset_results),indent=2))


if __name__=='__main__':main()