#!/usr/bin/env python3
"""Check exported weights, full covariances, precision loss and repeatability."""
import argparse
import json
from pathlib import Path
import numpy as np
from scipy.special import logsumexp
from rv_posterior import read_observations, transform_prior_coordinates, posterior_at_nodes


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--folder', type=Path, default=Path('.'))
    parser.add_argument('--out', type=Path, default=Path('posterior-checks.json'))
    arguments = parser.parse_args()
    folder = arguments.folder
    main_result = json.loads((folder / 'rv-posterior.json').read_text())
    repeated = json.loads((folder / 'rv-posterior-repeat.json').read_text())
    refined = json.loads((folder / 'rv-posterior-refined.json').read_text())
    output = []
    _, times, values, errors, _ = read_observations(folder)
    with np.load(folder / 'posterior-nodes.npz', allow_pickle=False) as archive:
        for result, repeat in zip(main_result['models'], repeated['models']):
            model = result['model']
            weights = archive[model + '_weights']
            density = archive[model + '_log_densities'] + archive[model + '_log_prior_minus_proposal']
            assert np.isclose(weights.sum(), 1., atol=1e-13)
            assert np.max(np.abs(weights - np.exp(density - logsumexp(density)))) < 1e-14
            assert 1 / np.sum(weights**2) > 3000
            means = archive[model + '_coefficient_means'].astype(float)
            dimension = means.shape[1]
            triangles = archive[model + '_coefficient_covariances_lower'].astype(float)
            covariances = np.zeros((len(weights), dimension, dimension))
            lower = np.tril_indices(dimension)
            covariances[:, lower[0], lower[1]] = triangles
            covariances[:, lower[1], lower[0]] = triangles
            assert np.linalg.eigvalsh(covariances).min() > 0
            selected = np.unique(np.r_[np.argsort(weights)[-100:], np.linspace(0, len(weights)-1, 32).astype(int)])
            nodes = transform_prior_coordinates(archive[model + '_prior_coordinates'][selected],
                main_result['priors']['rotation'], main_result['priors']['coupling_normal_scale'])
            recomputed = posterior_at_nodes(times, values, errors, nodes, model,
                main_result['priors']['planet_coefficient_normal_scale'])
            conditional_mean_error = float(np.max(np.abs(means[selected] - recomputed['coefficient_means'])))
            conditional_covariance_error = float(np.max(np.abs(covariances[selected] - recomputed['coefficient_covariances'])))
            conditional_mean_errors_by_axis = np.max(np.abs(means[selected] - recomputed['coefficient_means']), axis=0)
            conditional_covariance_errors_by_axis = np.max(np.abs(covariances[selected] - recomputed['coefficient_covariances']), axis=0)
            stored_mean = archive[model + '_coefficient_means']
            stored_covariance = archive[model + '_coefficient_covariances_lower']
            mean_rounding_bound = float(np.max(np.abs(np.spacing(stored_mean)))) / 2
            covariance_rounding_bound = float(np.max(np.abs(np.spacing(stored_covariance)))) / 2
            assert conditional_mean_error < 2e-5
            assert conditional_covariance_error < 2e-4
            mean = np.einsum('n,ni->i', weights, means)
            centered = means - mean
            covariance = np.einsum('n,nij->ij', weights, covariances) + np.einsum('n,ni,nj->ij', weights, centered, centered)
            mean_error = float(np.max(np.abs(mean - result['coefficient_mean'])))
            covariance_error = float(np.max(np.abs(covariance - result['coefficient_covariance'])))
            assert mean_error < 2e-6
            assert covariance_error < 2e-5
            global_covariance = np.array(result['joint_parameter_covariance'])
            assert np.max(np.abs(global_covariance - global_covariance.T)) < 1e-8
            assert np.linalg.eigvalsh(global_covariance).min() > 0
            differences = {}
            for label in ['b', 'c']:
                key = 'amplitude_' + label + '_m_s_quantiles_16_50_84'
                if key not in result:
                    continue
                differences[label] = float(np.max(np.abs(np.array(result[key]) - np.array(repeat[key]))))
                assert differences[label] < .08
            output.append(dict(model=model,normalized_weights=True,full_conditional_covariance_positive=True,
                full_joint_covariance_positive=True,export_float32_maximum_mean_error=mean_error,
                recomputed_conditional_nodes=len(selected),
                export_float32_maximum_selected_conditional_mean_error=conditional_mean_error,
                export_float32_maximum_selected_conditional_covariance_error=conditional_covariance_error,
                coefficient_axis_order=result['coefficient_names'],
                selected_conditional_mean_errors_by_axis=conditional_mean_errors_by_axis.tolist(),
                selected_conditional_covariance_errors_by_axis=conditional_covariance_errors_by_axis.tolist(),
                all_node_float32_round_to_nearest_absolute_mean_bound=mean_rounding_bound,
                all_node_float32_round_to_nearest_absolute_covariance_bound=covariance_rounding_bound,
                export_float32_maximum_covariance_error=covariance_error,repeat_maximum_amplitude_percentile_difference_m_s=differences))
    baseline = next(model for model in main_result['models'] if model['model'] == 'joint')
    fine = refined['models'][0]
    refinement = {}
    for label in ['b', 'c']:
        key = 'amplitude_' + label + '_m_s_quantiles_16_50_84'
        refinement[label] = float(np.max(np.abs(np.array(baseline[key]) - np.array(fine[key]))))
        assert refinement[label] < .08
    arguments.out.write_text(json.dumps(dict(models=output,refinement_maximum_amplitude_percentile_difference_m_s=refinement,
        scope='Export quantization and actual multi-seed/refined posterior comparisons under declared priors, not proof of astrophysical model completeness.'),indent=2)+'\n')
    print('AC18 posterior: exact weights, positive full covariances, float32 export and independent repeat/refinement checks pass')


if __name__ == '__main__':
    main()