"""Profile mass ratio and parallax amplitude; all nuisance parameters refit."""
from pathlib import Path
import argparse, json, platform, time
import numpy as np
from scipy.optimize import least_squares
from event_model import EventReference, compile_reference, load_observations, profile_magnitudes, projected_earth, 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('--parallax-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()))
library = folder / ('event-reference.dylib' if platform.system() == 'Darwin' else 'event-reference.so')
if not library.exists():
    library = compile_reference(folder)[0]
reference = EventReference(library)
earth = projected_earth(folder, observations['time'])
base = json.loads(arguments.baseline_results.read_text())
parallax = json.loads(arguments.parallax_results.read_text())
best = min((row for row in base['models'] if row['model'] == 'wide'), key=lambda row: row['chi_square'])
best_parallax = min((row for row in parallax['models'] if row['model'] == 'wide-annual-parallax' and row['impact_sign'] == 1), key=lambda row: row['chi_square'])
lower = np.array([3581.0, np.log(5), 0.05, 1.0001, np.log(1e-06), -np.pi, np.log(0.005), -np.pi])
upper = np.array([3584.5, np.log(60), 0.8, 2.5, np.log(0.03), np.pi, np.log(0.15), np.pi])
steps = np.array([0.0001, 0.0003, 0.0001, 0.0001, 0.003, 0.0001, 0.003, 0.003])
rows = []
for kind, values in [('mass_ratio', [1e-05, 2.5e-05, 4e-05, 6e-05, 0.0001, 0.0002, 0.0005]), ('parallax_amplitude', [0.0, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0])]:
    for fixed in values:
        if kind == 'mass_ratio':
            initial = np.array(best['parameters'])
            initial[4] = np.log(fixed)
            varying = [0, 1, 2, 3, 5, 6]
        else:
            initial = np.r_[best_parallax['parameters'][:7], np.arctan2(best_parallax['parameters'][8], best_parallax['parameters'][7])]
            varying = list(range(7)) if fixed == 0 else list(range(8))

        def physical(varying_values):
            parameters = initial.copy()
            parameters[varying] = varying_values
            if kind == 'parallax_amplitude':
                parameters = np.r_[parameters[:7], fixed * np.cos(parameters[7]), fixed * np.sin(parameters[7])]
            return parameters
        calls = [0]

        def residual(varying_values):
            parameters = physical(varying_values)
            source = source_trajectory(observations['time'], parameters, earth if kind == 'parallax_amplitude' else None)
            magnification = reference.binary(source, parameters[3], np.exp(parameters[4]), np.exp(parameters[6]), observations['gamma'])
            calls[0] += 1
            return profile_magnitudes(observations, magnification)[0]

        def jacobian(varying_values):
            columns = []
            for column, index in enumerate(varying):
                plus = varying_values.copy()
                minus = varying_values.copy()
                plus[column] = min(upper[index], varying_values[column] + steps[index])
                minus[column] = max(lower[index], varying_values[column] - steps[index])
                columns.append((residual(plus) - residual(minus)) / (plus[column] - minus[column]))
            return np.column_stack(columns)
        started = time.monotonic()
        fit = least_squares(residual, initial[varying], jac=jacobian, bounds=(lower[varying], upper[varying]), max_nfev=160, xtol=1e-10, ftol=1e-09, gtol=1e-06, x_scale='jac')
        errors = residual(fit.x)
        row = {'profile': kind, 'fixed_value': fixed, 'chi_square': float(errors @ errors), 'parameters': physical(fit.x).tolist(), 'optimizer_success': bool(fit.success), 'optimizer_message': fit.message, 'optimizer_evaluations': fit.nfev, 'actual_forward_calls': calls[0], 'active_bound_indices': [varying[i] for i in np.flatnonzero(fit.active_mask)], 'seconds': time.monotonic() - started}
        rows.append(row)
        (arguments.output / 'profile-history.json').write_text(json.dumps({'status': 'RUNNING', 'profiles': rows}, indent=2) + '\n')
        print(row, flush=True)
report = {'status': 'COMPLETED finite likelihood profiles; no posterior, coverage certificate or global-optimum claim', 'profiles': rows, 'conditioning': 'Positive-impact wide branch; supplied magnitude errors fixed; baseline and source fractions profiled separately for all five stations; every other named lens parameter refitted. Parallax amplitude profiles also fit direction.', 'bounds': {'lower': lower.tolist(), 'upper': upper.tolist()}, 'finite_difference_steps': steps.tolist()}
(arguments.output / 'profile-history.json').write_text(json.dumps(report, indent=2) + '\n')
print(report['status'], flush=True)