"""Fit the same 308 published measurements with explicitly bounded event models."""
from pathlib import Path
import argparse
import json
import time
import numpy as np
from scipy.optimize import least_squares
from event_model import EventReference, compile_reference, finite_point_lens, load_observations, profile_magnitudes, source_trajectory

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--materials', type=Path, required=True)
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--maximum-evaluations', type=int, default=100)
arguments = parser.parse_args()
folder = arguments.materials.resolve(); output = arguments.output.resolve(); output.mkdir(parents=True, exist_ok=True)
manifest = json.loads((folder / 'observation-sources.json').read_text())
observations = load_observations(folder / 'observations', manifest)
library, command = compile_reference(folder); reference = EventReference(library)
records = []
for model in ('point', 'wide', 'close'):
    if model == 'point':
        starts = [(1.61, 7.6e-5, -.3855926535897931)]
    elif model == 'wide':
        starts = [(1.61, ratio, -.3855926535897931) for ratio in (2e-5, 7.6e-5, 3e-4)] + [(1.03, .001, -.3855926535897931)]
    else:
        starts = [(separation, ratio, -.3855926535897931) for separation, ratio in ((.62, 7.6e-5), (.8, 3e-4), (.9, .001), (.98, .003))]
    for start_index, (separation, ratio, angle) in enumerate(starts):
        initial = np.array([3582.731, np.log(11.03), .359, separation, np.log(ratio), angle, np.log(.282 / 11.03)])
        varying = [0, 1, 2] if model == 'point' else list(range(7))
        lower = np.array([3581.0, np.log(5), .05, 1.0001 if model == 'wide' else .3, np.log(1e-6), -np.pi, np.log(.005)])
        upper = np.array([3584.5, np.log(60), .8, 2.5 if model == 'wide' else .9999, np.log(.03), np.pi, np.log(.15)])
        counter = [0]
        def residual(varying_values):
            parameters = initial.copy(); parameters[varying] = varying_values
            source = source_trajectory(observations['time'], parameters)
            magnification = (finite_point_lens(source, np.exp(parameters[6]), observations['gamma']) if model == 'point'
                             else reference.binary(source, parameters[3], np.exp(parameters[4]), np.exp(parameters[6]), observations['gamma']))
            counter[0] += 1
            return profile_magnitudes(observations, magnification)[0]
        differences = np.array([1e-4, 3e-4, 1e-4, 1e-4, .003, 1e-4, .003])[varying]
        def jacobian(values):
            columns = []
            for column, step in enumerate(differences):
                plus = values.copy(); minus = values.copy()
                plus[column] = min(upper[varying][column], values[column] + step)
                minus[column] = max(lower[varying][column], values[column] - step)
                columns.append((residual(plus) - residual(minus)) / (plus[column] - minus[column]))
            return np.column_stack(columns)
        started = time.monotonic()
        result = least_squares(residual, initial[varying], jac=jacobian, bounds=(lower[varying], upper[varying]),
                               max_nfev=arguments.maximum_evaluations, xtol=1e-10, ftol=1e-9, gtol=1e-6,
                               x_scale='jac')
        parameters = initial.copy(); parameters[varying] = result.x
        source = source_trajectory(observations['time'], parameters)
        magnification = (finite_point_lens(source, np.exp(parameters[6]), observations['gamma']) if model == 'point'
                         else reference.binary(source, parameters[3], np.exp(parameters[4]), np.exp(parameters[6]), observations['gamma']))
        errors, stations, predicted = profile_magnitudes(observations, magnification)
        row = {'model': model, 'start': start_index, 'initial': initial.tolist(), 'parameters': parameters.tolist(),
               'varying_lens_indices': varying, 'optimizer_success': bool(result.success), 'optimizer_message': result.message,
               'optimizer_evaluations': result.nfev, 'actual_forward_calls': counter[0], 'seconds': time.monotonic() - started,
               'chi_square': float(np.dot(errors, errors)), 'samples': len(errors), 'stations': stations,
               'active_bound_indices': [varying[index] for index, flag in enumerate(result.active_mask) if flag],
               'bounds': {'lower': lower.tolist(), 'upper': upper.tolist()},
               'finite_difference_steps': differences.tolist(),
               'maximum_absolute_standardized_residual': float(max(abs(errors))),
               'log_likelihood': float(-.5 * np.dot(errors, errors) - np.sum(np.log(observations['sigma'])) - len(errors) / 2 * np.log(2 * np.pi))}
        records.append(row)
        np.savetxt(output / f'{model}-{start_index}-local-predictions.csv', np.column_stack((observations['time'], observations['station'], observations['magnitude'], observations['sigma'], magnification, predicted, errors)), delimiter=',', header='original_time,station,observed_magnitude,reported_sigma,model_magnification,predicted_magnitude,standardized_residual', comments='')
        (output / 'fit-history.json').write_text(json.dumps({'status': 'RUNNING', 'models': records}, indent=2) + '\n')
        print(model, start_index, row['chi_square'], row['optimizer_success'], row['seconds'], row['parameters'], flush=True)
report = {'status': 'COMPLETED BOUNDED SEARCH, not a global-optimum certificate or posterior', 'samples': len(observations['time']),
          'likelihood': 'Independent Gaussian magnitudes, supplied errors unchanged, exact per-site baseline and fraction profiles',
          'point_source_radius_condition': .282 / 11.03, 'blend_source_fraction_domain': [0, 1],
          'reference_build': command, 'models': records,
          'local_output_notice': 'Per-row source-derived predictions stay in the user run; do not mirror them on the public site.'}
(output / 'fit-history.json').write_text(json.dumps(report, indent=2) + '\n')
print(report['status'], flush=True)