"""Compare zero-radius binary sources on the same original measured samples."""
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, load_observations, profile_magnitudes, source_trajectory
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--materials', type=Path, required=True)
parser.add_argument('--baseline-results', type=Path, required=True)
parser.add_argument('--output', type=Path, required=True)
arguments = parser.parse_args()
folder = arguments.materials.resolve()
arguments.output.mkdir(parents=True, exist_ok=True)
observations = load_observations(folder / 'observations', json.loads((folder / 'observation-sources.json').read_text()))
reference = EventReference(compile_reference(folder)[0])
base = json.loads(arguments.baseline_results.read_text())
records = []
for model in ('wide', 'close'):
    best = min((row for row in base['models'] if row['model'] == model), key=lambda row: row['chi_square'])
    for ratio_factor in (0.5, 1.0, 2.0):
        initial = np.array(best['parameters'])
        initial[4] += np.log(ratio_factor)
        lower = np.array([3581.0, np.log(5), 0.05, 1.0001 if model == 'wide' else 0.3, np.log(1e-06), -np.pi])
        upper = np.array([3584.5, np.log(60), 0.8, 2.5 if model == 'wide' else 0.9999, np.log(0.03), np.pi])
        steps = np.array([0.0001, 0.0003, 0.0001, 0.0001, 0.003, 0.0001])
        calls = [0]

        def residual(values):
            parameters = np.r_[values, initial[6]]
            source = source_trajectory(observations['time'], parameters)
            magnification = reference.binary(source, parameters[3], np.exp(parameters[4]), 0.0, observations['gamma'])
            calls[0] += 1
            return profile_magnitudes(observations, magnification)[0]

        def jacobian(values):
            columns = []
            for index, step in enumerate(steps):
                plus = values.copy()
                minus = values.copy()
                plus[index] = min(upper[index], values[index] + step)
                minus[index] = max(lower[index], values[index] - step)
                columns.append((residual(plus) - residual(minus)) / (plus[index] - minus[index]))
            return np.column_stack(columns)
        started = time.monotonic()
        fit = least_squares(residual, initial[:6], jac=jacobian, bounds=(lower, upper), max_nfev=200, xtol=1e-10, ftol=1e-09, gtol=1e-06, x_scale='jac')
        errors = residual(fit.x)
        row = {'model': model + '-point-source', 'fixed_source_radius': 0.0, 'initial_ratio_factor': ratio_factor, 'parameters': fit.x.tolist(), 'chi_square': float(errors @ errors), 'samples': len(errors), 'optimizer_success': bool(fit.success), 'optimizer_message': fit.message, 'optimizer_evaluations': fit.nfev, 'actual_forward_calls': calls[0], 'active_bound_indices': np.flatnonzero(fit.active_mask).tolist(), 'bounds': {'lower': lower.tolist(), 'upper': upper.tolist()}, 'seconds': time.monotonic() - started}
        records.append(row)
        (arguments.output / 'point-source-history.json').write_text(json.dumps({'status': 'RUNNING', 'models': records}, indent=2) + '\n')
        print(row, flush=True)
report = {'status': 'COMPLETED finite multi-start likelihood search; exact point-source instants, no source-radius clamp', 'models': records, 'exposure': 'Instantaneous because original exposure durations were not supplied; the same condition applies to the primary finite-source fit.'}
(arguments.output / 'point-source-history.json').write_text(json.dumps(report, indent=2) + '\n')
print(report['status'], flush=True)