"""Joint annual-parallax fits of the same five original photometry files."""
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, 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('--output', type=Path, required=True)
parser.add_argument('--maximum-evaluations', type=int, default=100)
parser.add_argument('--parallax-limit', type=float, default=10)
parser.add_argument('--seed-results', type=Path)
parser.add_argument('--only-model', choices=['wide','close'])
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)
base = json.loads(arguments.baseline_results.read_text()); library, command = compile_reference(folder); reference = EventReference(library)
earth = projected_earth(folder, observations['time'])
records = []
for model in ((arguments.only_model,) if arguments.only_model else ('wide', 'close')):
 best = min((row for row in base['models'] if row['model'] == model), key=lambda row: row['chi_square'])
 for sign in (1, -1):
  initial = np.r_[best['parameters'], 0., 0.]; initial[2] *= sign; initial[5] *= sign
  if arguments.seed_results:
   previous = json.loads(arguments.seed_results.read_text())
   seed = min((row for row in previous['models'] if row['model'] == model+'-annual-parallax' and row['impact_sign'] == sign), key=lambda row: row['chi_square'])
   initial = np.array(seed['parameters'])
  lower = np.array([3581., np.log(5), .05 if sign > 0 else -.8, 1.0001 if model == 'wide' else .3, np.log(1e-6), -np.pi, np.log(.005), -arguments.parallax_limit, -arguments.parallax_limit])
  upper = np.array([3584.5, np.log(60), .8 if sign > 0 else -.05, 2.5 if model == 'wide' else .9999, np.log(.03), np.pi, np.log(.15), arguments.parallax_limit, arguments.parallax_limit])
  differences = np.array([1e-4, 3e-4, 1e-4, 1e-4, .003, 1e-4, .003, .003, .003]); counter=[0]
  def residual(parameters):
   source = source_trajectory(observations['time'], parameters, earth)
   values = reference.binary(source, parameters[3], np.exp(parameters[4]), np.exp(parameters[6]), observations['gamma'])
   counter[0]+=1
   return profile_magnitudes(observations, values)[0]
  def jacobian(parameters):
   columns=[]
   for index,step in enumerate(differences):
    plus=parameters.copy();minus=parameters.copy();plus[index]=min(upper[index],parameters[index]+step);minus[index]=max(lower[index],parameters[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,jac=jacobian,bounds=(lower,upper),max_nfev=arguments.maximum_evaluations,xtol=1e-10,ftol=1e-9,gtol=1e-6,x_scale='jac')
  errors=residual(fit.x); row={'model':model+'-annual-parallax','impact_sign':sign,'initial':initial.tolist(),'parameters':fit.x.tolist(),
   'chi_square':float(errors@errors),'delta_chi_square_from_best_no_parallax':best['chi_square']-float(errors@errors),'samples':len(errors),
   'optimizer_success':bool(fit.success),'optimizer_message':fit.message,'optimizer_evaluations':fit.nfev,'actual_forward_calls':counter[0],
   'active_bound_indices':np.flatnonzero(fit.active_mask).tolist(),'bounds':{'lower':lower.tolist(),'upper':upper.tolist()},'finite_difference_steps':differences.tolist(),'seconds':time.monotonic()-started}
  records.append(row); (output/'parallax-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; not a posterior or a global optimum certificate','models':records,'reference_build':command,
 'time_qualification':'The original headerless photometry is evaluated on its quoted relative-day coordinate. DE441 TDB Earth vectors supply a conditional annual-parallax calculation; no exact HJD-to-BJD/TDB conversion is invented.',
 'convention':'pN,pE point along source-relative-to-lens motion; delta_tau=pN*delta_sN+pE*delta_sE, delta_beta=pN*delta_sE-pE*delta_sN; constant and linear Earth motion subtracted at reference day3582.731'}
(output/'parallax-history.json').write_text(json.dumps(report,indent=2)+'\n')
print(report['status'],flush=True)