## Run this script to produce the figures of Voisin et al. 2020
## Note that the modules "Fittriple.py" and "mcmcanalysis.py" must be in your pythonpath

import Fittriple as ft
import mcmcanalysis as mc
import numpy as np
import matplotlib.pyplot as pplot
import os
pplot.ioff()

fit = ft.Fittriple('parfile-reference-march2018-DE430-short-pset6-sepD', 'f4t1200s-FDM-20190514-gt600s-clean.tim-sorted-supershort')

start = 0
name = "data2019_10-SEPD_hires_1"
temp='1.00'
parfileres = "parfile-mean_SEPD_hires1_from0-pset6-sepD"
timfile = "4ft1200s-20191018-FDM-gt600s-superclean-no_sun_5deg.tim"


mcres = mc.mcmcresult(resfile='mcmcpt-' + name + '--temp='+temp+'.res',
                        parameter_object=fit, symbols=mc.symbols2)

# Do the chain plot
fig = mcres.Chain_plot(start=start, alpha = 0.5, average_window=288*12, show_initial=False, figsize=(16,100))
fig.savefig('figures/chain-'+ name + '-start' + '{:}'.format(start) + '-temp'+temp+'.png')

# Do the corner plot
# low res paper
fig = mcres.Corner_plot(start=start, show_tick_labels=False,
                    figsize=(20,20), rotationx=45, rotationy=30,
                     subplotadjust=(0.03,0.05, 0.99,0.99,0.,0.))
fig.savefig('figures/corner_plot.png')

# hi res supplematerial
fig = mcres.Corner_plot(start=start, show_tick_labels=True, figsize=None)
fig.savefig('figures/corner_plot-hires.png')


### Do the residuals

fit2 = ft.Fittriple(parfileres, timfile)
fig = fit2.Plot_time_residuals(mjd=True, show_both_wrms=False)
fig.savefig('figures/residuals.png')

fig = fit2.Plot_time_residuals_vs_orbitphases()
fig.savefig('figures/residuals_vs_orbphase.png')


### Do the histrograms
fontsize = 'x-large'
delta = np.abs(mcres.res[start:,mcres.inversparammap[20]])
delta.sort()
decil95 = delta[int(0.95 * delta.size)]

fig = pplot.figure(figsize=(10,5))
plt1 = fig.add_subplot(121)
plt2 = fig.add_subplot(122)
mcres.Plot_parameter_histogram(20,start=start, plt = plt1, fontsize=fontsize)

cnts, bins, patches  = plt2.hist(delta, bins=50)
plt2.axvline(decil95, linestyle='--', color='red', linewidth=1.5)
plt2.set_xlabel(r'$|\Delta|$', fontsize = fontsize)
plttop = plt2.twiny()
plttop.set_xlim(plt2.get_xlim())
plttop.set_xticks([decil95])
plttop.set_xticklabels(['{:.3g}'.format(decil95)], fontsize=fontsize)

plt2.ticklabel_format(style='sci', scilimits=(-1,1), axis='both')
arrowy = 0.1*cnts.max()
plt2.annotate("", xy=(0, arrowy), xytext=(decil95, arrowy), arrowprops=dict(arrowstyle="<->"))
plt2.text(0.5*decil95, 1.05*arrowy, "{:.2g}%".format(95), fontsize=fontsize)
plt2.tick_params(axis='both', which='major', labelsize=fontsize)
plt2.tick_params(axis='both', which='minor', labelsize=fontsize)

fig.set_tight_layout(True)
fig.show()
fig.savefig('figures/SEP_histograms.png')

### Make periodogram
fit2.Plot_Lombscargle_of_residuals(sampling_factor=0.2, extra_frequencies=[],
                                    show_low_peak=1, show_low_plus_earth=0,
                                    figsize=(12,6))
fit2.fig.savefig('figures/periodogram-paperversion.png')


### Check of accurracy :
print("\n\n   -------------------------  NOW RUNNING THE FOLLOWING COMMAND : --------------------------------------- ")
print('"Make_fake_timfile.exe parfile-mean_SEPD_hires1_from0-pset6-sepD 4ft1200s-20191018-FDM-gt600s-superclean-no_sun_5deg.tim-sorted 0. fake"')
print('MAKE SURE IT IS IN YOUR PATH !!')
print('\n\n')
os.system('Make_fake_timfile.exe parfile-mean_SEPD_hires1_from0-pset6-sepD 4ft1200s-20191018-FDM-gt600s-superclean-no_sun_5deg.tim-sorted 0. fake')
fit2 = ft.Fittriple(parfileres, 'fake')
fig =fit2.Plot_time_residuals()

import Fittriple as ft
fit = ft.Fittriple(parfileres, timfile)
fit.Compute_integrals_of_motion()
fig = fit.Plot_energies(fontsize='x-large', figsize=(8,4))
fig.show()
fig.savefig("figures/energy_conservation.png")

fig = fit.Plot_center_of_mass_positions(fit_quadratic_out=True, fontsize='x-large', figsize=(8,4))
fig.savefig('figures/centre_of_mass.png')
