"""Recompute own synthetic inference and aggregate fixed-source cluster checks."""
from pathlib import Path
import argparse,collections,json
import numpy as np
from scipy.stats import multivariate_normal
from scipy.linalg import cho_factor,cho_solve
from period_likelihood import HarmonicLikelihood
from frequency_measure import frequency_measure,period_posterior
from conditional_age import age_from_frequency_measure,age_from_reported_log_period
from empirical_calibration import EmpiricalCalibration
from inspect_cluster_samples import read_benchmarks,read_ruprecht
from age_calibration_prototype import fit_reference,reference_period
from rotation_reference import flux_series

def reproduce(inputs,output):
    output.mkdir(parents=True,exist_ok=True)
    table=inputs/'curtis2020-source/Benchmark_Clusters.txt';catalog=inputs/'curtis2020-source/Table_1-Ruprecht147_Catalog.txt'
    stars=read_benchmarks(table);old=read_ruprecht(catalog)
    assert len(stars)==923 and len(old)==440
    exact=[row for row in old if row['period_days']>0 and row['benchmark']=='Yes']
    assert len(exact)==35
    assert {row['source_id'] for row in exact}=={row['source_id'] for row in stars if row['cluster']=='Ruprecht 147'}
    calibration=EmpiricalCalibration(table)
    expected=json.loads(Path('calibration-aggregate.json').read_text())
    coefficient_error=max(abs(np.array(first['coefficients'])-np.array(second['coefficients'])).max() for first,second in zip(calibration.anchors,expected['anchors']))
    assert coefficient_error<1e-10
    native=json.loads(Path('synthetic-observations.json').read_text());rows=native['observations']
    times=np.array([row['timeDays'] for row in rows]);exposure=np.array([row['exposureDays'] for row in rows])
    likelihood=HarmonicLikelihood(times,[row['observedFlux'] for row in rows],[row['sigmaFlux'] for row in rows],exposure)
    evidence=likelihood.evaluate(native['frequencies']);evidence_error=float(max(abs(evidence-np.array(native['logEvidence']))))
    assert evidence_error<1e-6
    independent=[];cholesky_reference=[]
    for frequency in native['frequencies']:
        design=likelihood.design([frequency])[0];covariance=np.diag(likelihood.sigma**2)+(design*likelihood.prior_variances)@design.T
        independent.append(multivariate_normal.logpdf(likelihood.values,mean=np.zeros(len(times)),cov=covariance))
        factor=cho_factor(covariance,lower=True)
        cholesky_reference.append(-.5*(likelihood.values@cho_solve(factor,likelihood.values)+2*np.sum(np.log(np.diag(factor[0])))+len(times)*np.log(2*np.pi)))
    covariance_error=float(max(abs(evidence-independent)));cholesky_error=float(max(abs(evidence-cholesky_reference)));assert cholesky_error<2e-6 and covariance_error<2e-6
    measures=[];posteriors=[]
    for resolution in [16,32]:
        current=HarmonicLikelihood(times,[row['observedFlux'] for row in rows],.0005,exposure)
        measure=frequency_measure(current,samples_per_cycle=resolution);assert measure['status']=='integrated',measure
        measures.append(measure);posteriors.append(period_posterior(measure,current))
    mass_error=abs(np.sum(measures[0]['scaled_likelihood_log_frequency_mass'])*np.exp(measures[0]['log_scale']-measures[1]['log_scale'])/np.sum(measures[1]['scaled_likelihood_log_frequency_mass'])-1)
    assert mass_error<.001
    age=age_from_frequency_measure(calibration,measures[0],5500,40,age_points=401,temperature_order=32)
    refined=age_from_frequency_measure(calibration,measures[1],5500,40,age_points=801,temperature_order=64)
    variation=float(.5*np.trapezoid(abs(np.interp(refined['age_gyr'],age['age_gyr'],age['density_per_gyr'])-refined['density_per_gyr']),refined['age_gyr']))
    assert variation<.001 and abs(np.trapezoid(age['density_per_gyr'],age['age_gyr'])-1)<1e-12
    nodes,weights=np.polynomial.legendre.leggauss(17);flux=np.zeros(len(times))
    for node,weight in zip(nodes,weights):flux+=weight/2*flux_series(times+exposure*node/2,order=(96,480))[0]
    flux_error=float(max(abs(flux-native['forwardFlux'])));assert flux_error<1e-6
    training=[row for row in stars if row['cluster']=='Praesepe' and 4200<=row['teff_k']<=6000];coefficients=fit_reference(training)
    failures=[]
    for low,high in [(4200,4600),(5500,6000)]:
        selected=[row for row in stars if row['cluster']=='NGC 6811' and low<=row['teff_k']<high]
        for exponent in [.5,.62]:
            inferred=.67*(np.array([row['period_days'] for row in selected])/reference_period(coefficients,[row['teff_k'] for row in selected]))**(1/exponent)
            failures.append({'temperature_bin_k':[low,high],'stars':len(selected),'imposed_exponent':exponent,'median_conditional_age_gyr':float(np.median(inferred))})
    holdouts=[]
    for cluster,model in [('NGC 6819',calibration),('NGC 752',calibration),('NGC 6811',EmpiricalCalibration(table,excluded_cluster='NGC 6811'))]:
        selected=[row for row in stars if row['cluster']==cluster and model.temperature_bounds[0]<=row['teff_k']<=model.temperature_bounds[1]]
        summaries=[]
        for row in selected:
            posterior=age_from_reported_log_period(model,np.log(row['period_days']),0.,row['teff_k'],sigma_temperature=0.)
            q=posterior['quantiles025_16_50_84_975_gyr'];summaries.append([q[2],q[4]-q[0],q[0]<=row['age_gyr']<=q[4],bool(posterior['flags'])])
        values=np.array(summaries)
        holdouts.append({'cluster':cluster,'stars':len(selected),'fixed_source_age_gyr':selected[0]['age_gyr'],'median_conditional_age_gyr':float(np.median(values[:,0])),'median95_width_gyr':float(np.median(values[:,1])),'fraction_including_fixed_age':float(np.mean(values[:,2])),'fraction_boundary_notices':float(np.mean(values[:,3]))})
    report={'status':'PASS independently recomputed own synthetic likelihood/flux/age and aggregate original-source cluster comparison','benchmark_rows':len(stars),'catalog_rows':len(old),'exact_yes_period_positive':len(exact),'cluster_counts':dict(collections.Counter(row['cluster'] for row in stars)),
            'coefficient_error':float(coefficient_error),'log_evidence_error_vs_JS':evidence_error,'log_evidence_error_vs_full_covariance_eigen':covariance_error,'log_evidence_error_vs_full_covariance_cholesky':cholesky_error,'full_covariance_absolute_log_tolerance':2e-6,'relative_frequency_mass_refinement':float(mass_error),'age_density_refinement_TV':variation,'independent_flux_max_error':flux_error,
            'default_age_quantiles_gyr':age['quantiles025_16_50_84_975_gyr'].tolist(),'default_clock_hypothesis_probabilities':posteriors[0]['harmonic_model_probabilities'],'separable_failures':failures,'whole_cluster_holdouts':holdouts,
            'qualification':'Fixed author temperature/membership and0.67/1/2.7Gyr scale. Catalog periods/temperatures condition exactly because no homogeneous measurement errors supplied; effective predictive width is not measured intrinsic scatter. No universal coverage or absolute age precision.'}
    (output/'summary.json').write_text(json.dumps(report,indent=2)+'\n')
    # Only self-authored synthetic density is exported point by point.
    (output/'synthetic-age-density.json').write_text(json.dumps({'age_gyr':age['age_gyr'].tolist(),'density_per_gyr':age['density_per_gyr'].tolist(),'kind':'Own synthetic saved-data conditional result'})+'\n')
    print(json.dumps(report,indent=2))
if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--inputs',type=Path,default=Path('inputs'));parser.add_argument('--output',type=Path,default=Path('results'));args=parser.parse_args();reproduce(args.inputs,args.output)