#!/usr/bin/env python3
"""Conditional latent RV/calcium posterior at the71retained measurement times.

This is an in-sample model distribution, not new observations or an error-
coverage test. Whole time/channel covariance is saved, not just diagonal bars.
"""
import argparse,json
from pathlib import Path
import numpy as np
from scipy.linalg import cho_factor,cho_solve
from rv_kernel import activity_covariance
from rv_posterior import read_observations,transform_prior_coordinates,design_matrix,HYPER_NAMES


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--data',type=Path,default=Path('.'))
    parser.add_argument('--posterior',type=Path,default=Path('rv-posterior.json'))
    parser.add_argument('--nodes',type=Path,default=Path('posterior-nodes.npz'))
    parser.add_argument('--draws',type=int,default=1024);parser.add_argument('--out',type=Path,default=Path('posterior-predictions.json'))
    args=parser.parse_args();posterior=json.loads(args.posterior.read_text())
    rows,times,values,errors,normalization=read_observations(args.data);count=len(times)
    random=np.random.default_rng(1819);outputs=[]
    with np.load(args.nodes,allow_pickle=False) as archive:
        for model in ['kepler','activity','joint']:
            coordinates=archive[model+'_prior_coordinates'];weights=archive[model+'_weights'];weights=weights/weights.sum()
            indices=np.searchsorted(np.cumsum(weights),(np.arange(args.draws)+.5)/args.draws)
            nodes=transform_prior_coordinates(coordinates[indices],posterior['priors']['rotation'],posterior['priors']['coupling_normal_scale'])
            means=[];covariances=[];draws=[];activity_means=[];minimum_eigenvalue=0.
            for node in nodes:
                parameters=dict(zip(HYPER_NAMES,node))
                activity=np.zeros((2*count,2*count)) if model=='kepler' else activity_covariance(times,times,parameters)
                planets=model!='activity';design=design_matrix(times,node[8],planets)
                scales=np.array([100.,10.]+([posterior['priors']['planet_coefficient_normal_scale']]*4 if planets else []))
                scaled_design=design*scales**2
                latent=activity+np.einsum('ik,jk->ij',scaled_design,design)
                noise=errors**2+np.r_[np.full(count,node[6]**2),np.full(count,node[7]**2)]
                total=latent+np.diag(noise)
                factor=cho_factor(total,lower=True,check_finite=False)
                inverse_values=cho_solve(factor,values,check_finite=False)
                inverse=cho_solve(factor,np.eye(2*count),check_finite=False)
                mean=values-noise*inverse_values
                covariance=np.diag(noise)-noise[:,None]*inverse*noise[None,:]
                covariance=(covariance+covariance.T)/2
                eigenvalues,eigenvectors=np.linalg.eigh(covariance)
                minimum_eigenvalue=min(minimum_eigenvalue,float(eigenvalues.min()))
                if eigenvalues.min() < -1e-6:raise ValueError('conditional latent covariance failed positivity')
                draw=mean+np.einsum('ij,j->i',eigenvectors,np.sqrt(np.maximum(0,eigenvalues))*random.normal(size=2*count))
                means.append(mean);covariances.append(covariance);draws.append(draw)
                activity_means.append(np.einsum('ij,j->i',activity,inverse_values))
            means=np.array(means);draws=np.array(draws)
            mean=means.mean(axis=0);centered=means-mean
            covariance=np.mean(covariances,axis=0)+np.einsum('ni,nj->ij',centered,centered)/len(means)
            quantiles=np.quantile(draws,[.16,.5,.84],axis=0)
            activity_mean=np.mean(activity_means,axis=0)
            samples=[]
            for index,row in enumerate(rows):
                samples.append(dict(time_bjd_utc=float(row['bjd_utc_day']),time_offset_days=float(times[index]),
                    observed_rv_m_s=float(values[index]),source_rv_error_m_s=float(errors[index]),
                    observed_log_rhk=float(row['log_rhk']),source_log_rhk_error=float(row['log_rhk_error']),
                    posterior_rv_mean_m_s=float(mean[index]),posterior_rv_quantiles_16_50_84=quantiles[:,index].tolist(),
                    posterior_calcium_mean_log_rhk=float(mean[count+index]*.01+normalization['calcium_origin_log_rhk']),
                    posterior_calcium_quantiles_16_50_84=(quantiles[:,count+index]*.01+normalization['calcium_origin_log_rhk']).tolist(),
                    activity_rv_mean_m_s=float(activity_mean[index])))
            outputs.append(dict(model=model,samples=samples,full_latent_mean=mean.tolist(),full_latent_covariance=covariance.tolist(),
                covariance_axis_order='first71RVmetres/second, then71calcium centered/dividedby0.01dex',
                conditional_covariance_min_eigenvalue_before_roundoff_correction=minimum_eigenvalue,
                normalized_rv_residual_rms=float(np.sqrt(np.mean((values[:count]-mean[:count])**2)))))
    result=dict(schema='ac18-in-sample-posterior-v1',kind='model latent posterior conditional on the real retained data, not measured additional signals',
        posterior_mixture_draws=args.draws,hyperparameter_selection='deterministic stratified resampling of all retained normalized posterior weights',
        seed=1819,normalization=normalization,models=outputs,
        interpretation='Pointwise16/84percentiles include model uncertainty under the declared priors. They are neither new observational errors nor calibrated false-positive probabilities.')
    args.out.write_text(json.dumps(result,indent=2,allow_nan=False)+'\n')
    print('AC18: full71-by2 latent covariance and posterior summaries from',args.draws,'mixture draws')


if __name__=='__main__':main()