"""Recompute frozen model-level statistics from origin-downloaded observations.

This forward replay is independent of the website and is not an optimizer rerun.
Only aggregate statistics are written; local fit scripts can separately retain
row-wise predictions in your own working folder.
"""
from pathlib import Path
import argparse
import json
import numpy as np
from event_model import EventReference, compile_reference, finite_point_lens, load_observations, profile_magnitudes, projected_earth, source_trajectory


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--materials', type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument('--output', type=Path, default=Path('replay-summary.json'))
    arguments = parser.parse_args()
    folder = arguments.materials.resolve()
    observations = load_observations(folder / 'observations', json.loads((folder / 'observation-sources.json').read_text()))
    reference = EventReference(compile_reference(folder)[0])
    aggregate = json.loads((folder / 'aggregate-results.json').read_text())
    earth = projected_earth(folder, observations['time'])
    reports = []
    for name, report in aggregate['reports'].items():
        candidates = report.get('models', report.get('profiles', []))
        for index, candidate in enumerate(candidates):
            parameters = np.array(candidate['parameters'])
            if len(parameters) == 6:
                parameters = np.r_[parameters, np.log(.01)]  # Radius is explicitly zero in this model, not read from this unused slot.
            source = source_trajectory(observations['time'], parameters, earth if len(parameters) == 9 else None)
            model = candidate.get('model', candidate.get('profile'))
            radius = 0 if name == 'zero_radius_binary' else np.exp(parameters[6])
            magnification = finite_point_lens(source, radius, observations['gamma']) if model == 'point' else reference.binary(source, parameters[3], np.exp(parameters[4]), radius, observations['gamma'])
            residual = profile_magnitudes(observations, magnification)[0]
            value = float(residual @ residual)
            difference = abs(value - candidate['chi_square'])
            if difference > .15:
                raise AssertionError((name, index, model, value, candidate['chi_square'], difference))
            reports.append({'report': name, 'index': index, 'model': model, 'chi_square': value, 'absolute_replay_difference': difference})
    result = {'status': 'PASS frozen forward replay; independent local optimization is a separate command',
              'samples': len(observations['time']), 'tolerance_absolute_chi_square': .15,
              'maximum_absolute_replay_difference': max(row['absolute_replay_difference'] for row in reports), 'models_and_profiles': reports}
    arguments.output.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps({key: value for key, value in result.items() if key != 'models_and_profiles'}, indent=2))


if __name__ == '__main__': main()