#!/usr/bin/env python3
"""AC18 standalone real RV/activity posterior; public authored science, MIT."""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
import scipy
from rv_posterior import prior_nodes,read_observations,posterior_at_nodes,summarize,HYPER_NAMES


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument("--data",type=Path,default=Path("."))
    parser.add_argument("--power",type=int,default=11)
    parser.add_argument("--seed",type=int,default=1801)
    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-posterior.json"))
    parser.add_argument("--nodes-out",type=Path)
    args=parser.parse_args()
    if not 7<=args.power<=17: raise ValueError("power must be between7and17")
    rows,times,values,errors,normalization=read_observations(args.data)
    nodes=prior_nodes(args.power,args.seed,args.rotation_prior,args.amplitude_scale)
    result=dict(schema="ac18-rv-posterior-v1",kind="real published RV/calcium observations under a finite proper hyperparameter prior and continuous Gaussian coefficient priors",
        numpy_version=np.__version__,scipy_version=scipy.__version__,source_rows=len(rows),
        input_sha256=hashlib.sha256((args.data/"harps2012-corot7.csv").read_bytes()).hexdigest(),
        normalization=normalization,assumptions=["circular orbits; first period conditioned on published0.85359165day value; phases free",
        "second period finite prior within3.3to4.1days; not a blind frequency search",
        "RV and normalized measured logRHK jointly modeled; FWHM/BIS retained ancillary only, with no invented point errors",
        "G/Gdot Gaussian local statistical model, distinct from browser surface simulation",
        "all normalized finite-prior weights and full coefficient/joint covariance retained; no peak likelihood model probability"],
        prior=dict(nodes=len(nodes),equal_prior_mass=1/len(nodes),sobol_seed=args.seed,
            rotation=args.rotation_prior,rv_coupling_normal_scale_m_s=args.amplitude_scale,
            derivative_coupling_normal_scale_m_s_day=5*args.amplitude_scale,
            planet_sine_cosine_normal_scale_m_s=args.planet_scale,
            orbit_amplitude_prior="Rayleigh implied by independent zero-mean Gaussian sine/cosine coefficients",
            hyper_names=HYPER_NAMES),models=[],samples=[dict(time_bjd_utc=float(row["bjd_utc_day"]),
            rv_m_s=float(row["rv_km_s"])*1000,rv_error_m_s=float(row["rv_error_km_s"])*1000,
            log_rhk=float(row["log_rhk"]),log_rhk_error=float(row["log_rhk_error"]),
            fwhm_m_s=float(row["fwhm_km_s"])*1000,bis_m_s=float(row["bis_km_s"])*1000) for row in rows])
    frozen={"prior_nodes":nodes}
    for model in ["kepler","activity","joint"]:
        posterior=posterior_at_nodes(times,values,errors,nodes,model,args.planet_scale)
        result["models"].append(summarize(posterior,nodes))
        for key in ["weights","log_densities","coefficient_means","coefficient_covariances"]:
            frozen[model+"_"+key]=posterior[key]
    args.out.parent.mkdir(parents=True,exist_ok=True)
    args.out.write_text(json.dumps(result,indent=2,allow_nan=False)+"\n")
    if args.nodes_out:
        np.savez_compressed(args.nodes_out,**frozen)
    print(json.dumps([dict(model=m["model"],effective_nodes=m["effective_nodes"],max_weight=m["maximum_node_weight"],
        amplitude_b=m.get("amplitude_b_m_s_quantiles_16_50_84"),amplitude_c=m.get("amplitude_c_m_s_quantiles_16_50_84")) for m in result["models"]],indent=2))


if __name__=="__main__":
    main()