#!/usr/bin/env python3
"""Continuous-prior defensive multi-center RV/activity importance posterior."""
import argparse,json,time
from pathlib import Path
import numpy as np
from rv_posterior import read_observations,transform_prior_coordinates,posterior_at_nodes,summarize
from rv_importance import fit_defensive_proposal,standard_normal_log_density,importance_diagnostics,adapt_defensive_proposal


def run_model(times,values,errors,model,power,seed,rotation_prior,amplitude_scale,planet_scale):
    random=np.random.default_rng(seed)
    def target(coordinates):
        node=transform_prior_coordinates(coordinates,rotation_prior,amplitude_scale)
        posterior=posterior_at_nodes(times,values,errors,node,model,planet_scale)
        value=posterior['log_densities'][0]+standard_normal_log_density(coordinates)
        return float(value) if np.isfinite(value) else -1e100
    exploration=[]
    for rotation in [-2.,0.,2.]:
        for orbital in [-2.,0.,2.]:
            point=np.zeros(9);point[1]=rotation;point[8]=orbital;exploration.append(point)
    proposal,fits=fit_defensive_proposal(target,9,random,probe_count=512,maximum_modes=8,
                                       exploration_starts=exploration)
    initial_proposal=proposal.metadata()
    pilot_diagnostics=[]
    for pilot_count in [32768,16384]:
        pilot=proposal.sample(pilot_count,random)
        pilot_nodes=transform_prior_coordinates(pilot,rotation_prior,amplitude_scale)
        pilot_correction=standard_normal_log_density(pilot)-proposal.log_density(pilot)
        pilot_posterior=posterior_at_nodes(times,values,errors,pilot_nodes,model,planet_scale,pilot_correction)
        pilot_diagnostics.append(dict(samples=pilot_count,**importance_diagnostics(pilot_posterior['log_densities']+pilot_correction)))
        proposal=adapt_defensive_proposal(pilot,pilot_posterior['log_densities']+pilot_correction,random)
    coordinates=proposal.sample(2**power,random)
    nodes=transform_prior_coordinates(coordinates,rotation_prior,amplitude_scale)
    correction=standard_normal_log_density(coordinates)-proposal.log_density(coordinates)
    posterior=posterior_at_nodes(times,values,errors,nodes,model,planet_scale,correction)
    summary=summarize(posterior,nodes)
    summary['continuous_prior_log_marginal_estimate']=summary.pop('finite_prior_log_marginal')
    summary['importance']=importance_diagnostics(posterior['log_densities']+correction)
    summary['proposal']=proposal.metadata();summary['initial_proposal']=initial_proposal
    summary['pilot_diagnostics']=pilot_diagnostics;summary['proposal_fits']=fits
    summary['prior_coordinate_seed']=seed;summary['samples']=len(nodes)
    return summary,dict(nodes=nodes,prior_coordinates=coordinates,log_prior_minus_proposal=correction,
        weights=posterior['weights'],log_densities=posterior['log_densities'],
        coefficient_means=posterior['coefficient_means'],coefficient_covariances=posterior['coefficient_covariances'])


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--data',type=Path,default=Path('.'))
    parser.add_argument('--power',type=int,default=12);parser.add_argument('--seed',type=int,default=1808)
    parser.add_argument('--model',choices=['kepler','activity','joint','all'],default='all')
    parser.add_argument('--rotation-prior',choices=['broad','photometry','wide'],default='broad')
    parser.add_argument('--amplitude-scale',type=float,default=10.);parser.add_argument('--planet-scale',type=float,default=10.)
    parser.add_argument('--out',type=Path,default=Path('rv-importance.json'));parser.add_argument('--nodes-out',type=Path)
    args=parser.parse_args();rows,times,values,errors,normalization=read_observations(args.data)
    output=dict(schema='ac18-rv-importance-v1',kind='real RV/activity continuous proper-prior posterior approximated by defensive multi-center importance sampling',
        normalization=normalization,seed=args.seed,power=args.power,source_rows=len(rows),
        assumptions=['fixed declared offset-prior references31180metres/second and minus4.72logRHK', 'independent quoted channel errors plus free scatter; source supplies no RV/calcium error covariance', 'instantaneous evaluation at published BJD_UTC; no unprovided exposure kernel is invented'],
        priors=dict(rotation=args.rotation_prior,coupling_normal_scale=args.amplitude_scale,planet_coefficient_normal_scale=args.planet_scale),models=[])
    frozen={};start=time.perf_counter()
    for index,model in enumerate(['kepler','activity','joint'] if args.model=='all' else [args.model]):
        summary,arrays=run_model(times,values,errors,model,args.power,args.seed+index,args.rotation_prior,args.amplitude_scale,args.planet_scale)
        output['models'].append(summary)
        for name,array in arrays.items():frozen[model+'_'+name]=array
        print(json.dumps(dict(model=model,seconds=time.perf_counter()-start,diagnostic=summary['importance'],
            amplitude_b=summary.get('amplitude_b_m_s_quantiles_16_50_84'),amplitude_c=summary.get('amplitude_c_m_s_quantiles_16_50_84')),indent=2),flush=True)
    args.out.parent.mkdir(exist_ok=True,parents=True);args.out.write_text(json.dumps(output,indent=2,allow_nan=False)+'\n')
    if args.nodes_out:
        compact={}
        for name,array in frozen.items():
            if name.endswith('_nodes'):continue  # regenerated exactly from saved prior coordinates
            if name.endswith('_coefficient_covariances'):
                row,column=np.tril_indices(array.shape[1])
                compact[name+'_lower']=array[:,row,column].astype(np.float32)
            elif name.endswith('_coefficient_means'):compact[name]=array.astype(np.float32)
            else:compact[name]=array
        compact['storage_version']=np.array([1],dtype=np.int32)
        output['node_storage']={'prior_coordinates_weights_densities':'float64',
            'coefficient_means_and_full_symmetric_covariance_lower_triangle':'float32 export; inference and global covariance computed in float64',
            'nodes':'regenerate physical hyperparameters from exact stored Gaussian-prior coordinates'}
        args.out.write_text(json.dumps(output,indent=2,allow_nan=False)+'\n')
        np.savez_compressed(args.nodes_out,**compact)


if __name__=='__main__':main()