from backtracks import System
import multiprocessing as mp

LOAD = False

# target = "HD 131399A"
# target = "Kepler-1900"
# file = "scorpions1b_orbitizelike.csv"
# sample_method = 'unif'
# yy_epochs = False

# target = "YSES2"
# file = "yses2b.csv"
# file = "yses2b-orbitizelike-latest.csv"
# sample_method = 'rslice'
# yy_epochs = False

# target = "SAO206463"
# file = "sao206463b_astrometry_local.csv"
# sample_method = 'rwalk'
# yy_epochs = False

target = "HD165054"
file = "hd165054_0_orbitizelike.csv"
sample_method = 'rwalk'
yy_epochs = True


track = System(target_name=target,
               candidate_file=file,
               nearby_window=0.5,
               relax_pm_priors=False,
               relax_par_priors=False,
               yy_epochs=yy_epochs,
               fileprefix='./')


fig_stat=track.generate_stationary_plot(days_backward=1*365,
                                         days_forward=7.*365,
                                         step_size=10.,
                                         plot_radec=False,
                                         fileprefix='./',
                                         filepost='.pdf')


if LOAD:
    track.load_results(fileprefix='./')
else:
    results = track.fit(dlogz=1e-1, npool=mp.cpu_count(), 
                        dynamic=True, nlive=100,
                        mpi_pool=False, 
                        resume=False, 
                        sample_method=sample_method)
    track.save_results(fileprefix='./')


fig_track, fig_post, fig_diag, fig_prior, fig_hood = \
    track.generate_plots(days_backward=1.*365.,
                         days_forward=7.*365,
                         step_size=10.,
                         plot_radec=False,
                         plot_stationary=False,
                         fileprefix='./',
                         filepost='.pdf'
                         )


# import numpy as np
# from astropy.time import Time
# from backtracks.utils import utc2tt

# import orbitize
# import orbitize.driver
# from orbitize import priors
# import matplotlib.pyplot as plt

# my_driver = orbitize.driver.Driver('./scorpions1b_orbitizelike_test.csv', 
#                                    'OFTI',
#                                    2, # number of secondary bodies in system
#                                    1.10, # total mass [M_sun]
#                                    9.1115, # total parallax of system [mas]
#                                    mass_err=0.03,
#                                    plx_err=0.0310)


# sys = my_driver.system
# lab = sys.param_idx
# print(sys.labels)

# sys.sys_priors[lab['sma2']] = priors.UniformPrior(0.1, 10)


# s = my_driver.sampler
# orbits = s.run_sampler(100)

# myResults = s.results

# epochs = my_driver.system.data_table['epoch']
# orbit_figure = myResults.plot_orbits(
#     num_orbits_to_plot=15,
#     start_mjd=epochs[0], # Minimum MJD for colorbar (here we choose first data epoch)
#     cmap=plt.cm.magma,
#     cbar_param='ecc1',
#     show_colorbar=False,
#     plot_astrometry=False,
#     sep_pa_color='xkcd:purple',
#     fig=fig_track
# )

# ax_orb, ax_sep, ax_pa  = orbit_figure.axes

# # ax_orb.set_xlim(-200,-100)
# # ax_orb.set_ylim(-40,60)

# # create an array of times
# projected_epochs = np.arange(Time('2024-12-15').jd, Time('2025-05-01').jd, 10)
# projected_epochs_tt = utc2tt(projected_epochs)

# # we need an array of parameters, which are
# # ra, dec, pmra, pmdec, par, ra_host, dec_host, pmra_host, pmdec_host, par_host, rv_host
# # in many cases, we're just interested in the stationary case, so we can pull the array of those parameters directly from our System
# parameters = track.run_median

# # now, we just call the System.radecdists function
# ras, decs = track.radecdists(projected_epochs_tt, parameters) # these array have units mas

# # and we'll plot the results to prove this works
# # plt.figure(figsize=(9,6))
# ax_orb.plot(ras, decs, c='cornflowerblue', zorder=1)

# projected_epochs = np.arange(Time('2018-10-01').jd, Time('2025-05-01').jd, 10)
# projected_epochs_tt = utc2tt(projected_epochs)
# ras, decs = track.radecdists(projected_epochs_tt, parameters)
# ax_orb.plot(ras, decs, c='k', zorder=-1, ls='--')

# plt.xlim(-380, -470)
# plt.ylim(-960, -930)

# # and we can throw up some data
# ax_orb.errorbar(track.ras, track.decs, yerr=track.decserr, xerr=track.raserr,
#              color="tab:gray", ecolor='tomato', mec='tomato',
#              linestyle="none", marker="o", ms=5., mew=1.5)

# # and label the axes
# # plt.xlabel('$\Delta$RA (mas)')
# # plt.ylabel('$\Delta$Dec (mas)')

# plt.savefig('test.pdf', dpi=300, bbox_inches='tight')

# print('median track chi2_red is',track.median_chi2_red)
# print('stationary track chi2_red is',track.stationary_chi2_red)
# print('dlogz between stationary and median is',track.stationary_loglike-track.median_loglike)
