"""Continue failed profiles and seek a better local minimum from adjacent solutions.

The first pass is retained unchanged. This finite search is not a proof that
all disconnected solutions have been found.
"""
from pathlib import Path
import argparse
import json
import platform
import time
import numpy as np
from scipy.optimize import least_squares
from event_model import EventReference, load_observations, profile_magnitudes, projected_earth, source_trajectory


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--materials', type=Path, required=True)
    parser.add_argument('--profiles', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--maximum-evaluations', type=int, default=500)
    parser.add_argument('--only-interpolated-parallax', action='store_true')
    arguments = parser.parse_args()
    folder = arguments.materials.resolve()
    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')
    reference = EventReference(library)
    earth = projected_earth(folder, observations['time'])
    first = json.loads(arguments.profiles.read_text())
    lower = np.array([3581., np.log(5), .05, 1.0001, np.log(1e-6), -np.pi, np.log(.005), -np.pi])
    upper = np.array([3584.5, np.log(60), .8, 2.5, np.log(.03), np.pi, np.log(.15), np.pi])
    steps = np.array([1e-4, 3e-4, 1e-4, 1e-4, .003, 1e-4, .003, .003])
    attempts = []
    targets = [('mass_ratio', 1e-5, [1e-5, 2.5e-5]), ('parallax_amplitude', 5., [2., 5., 10., 'interpolated2to10'])]
    if arguments.only_interpolated_parallax:
        targets = [('parallax_amplitude', 5., ['interpolated2to10'])]
    for kind, fixed, seed_values in targets:
        for seed_value in seed_values:
            if seed_value == 'interpolated2to10':
                neighbors = [next(row for row in first['profiles'] if row['profile'] == kind and row['fixed_value'] == value) for value in [2., 10.]]
                original = np.array(neighbors[0]['parameters']) + 3 / 8 * (np.array(neighbors[1]['parameters']) - np.array(neighbors[0]['parameters']))
            else:
                seed = next(row for row in first['profiles'] if row['profile'] == kind and row['fixed_value'] == seed_value)
                original = np.array(seed['parameters'])
            if kind == 'mass_ratio':
                initial = original.copy()
                initial[4] = np.log(fixed)
                varying = [0, 1, 2, 3, 5, 6]
            else:
                initial = np.r_[original[:7], np.arctan2(original[8], original[7])]
                varying = list(range(8))
            def physical(values):
                parameters = initial.copy()
                parameters[varying] = values
                return np.r_[parameters[:7], fixed * np.cos(parameters[7]), fixed * np.sin(parameters[7])] if kind == 'parallax_amplitude' else parameters
            calls = [0]
            def residual(values):
                parameters = physical(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(values):
                columns = []
                for column, index in enumerate(varying):
                    plus = values.copy(); minus = values.copy()
                    plus[column] = min(upper[index], values[column] + steps[index])
                    minus[column] = max(lower[index], 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=arguments.maximum_evaluations, xtol=1e-10, ftol=1e-9, gtol=1e-6, x_scale='jac')
            errors = residual(fit.x)
            row = {'profile': kind, 'fixed_value': fixed, 'seed_fixed_value': seed_value,
                   'parameters': physical(fit.x).tolist(), 'chi_square': float(errors @ errors),
                   'optimizer_success': bool(fit.success), 'optimizer_message': fit.message,
                   'optimizer_evaluations': fit.nfev, 'actual_forward_calls': calls[0],
                   'seconds': time.monotonic() - started}
            attempts.append(row)
            arguments.output.mkdir(parents=True, exist_ok=True)
            (arguments.output / 'refinement-attempts.json').write_text(json.dumps(attempts, indent=2) + '\n')
            print(row, flush=True)
    selected = []
    for previous in first['profiles']:
        choices = [previous] + [row for row in attempts if row['profile'] == previous['profile'] and row['fixed_value'] == previous['fixed_value']]
        successful = [row for row in choices if row['optimizer_success']]
        selected.append(min(successful or choices, key=lambda row: row['chi_square']))
    report = {**first, 'status': 'COMPLETED finite best-found profiles; adjacent-seed refinement retained separately; no global or interval-coverage certificate',
              'profiles': selected, 'refinement_attempts': first.get('refinement_attempts', []) + attempts}
    (arguments.output / 'profile-history.json').write_text(json.dumps(report, indent=2) + '\n')


if __name__ == '__main__':
    main()