import at
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import factorial

#simple_ring = at.simple_ring(3e9, 528, 176, 0.1, 0.1, 1e6, 0.000306, U0=0.0, tauz=1000, espread=0.000769,Qpy=1.0, chromx_arr=[0,1e3,1e5])
simple_ring = at.simple_ring(3e9, 528, 176, 0.1, 0.1, 1e6, 0.000306, U0=0.0, tauz=1000, espread=0.000769,Qpy=1.0)
simple_ring.enable_6d()

fitOrder = 4
(fitH, fitV), dpa, qz = at.chromaticity(simple_ring, method='linopt', dpm=3e-3, npoints=20, order=fitOrder, dp=0)

fitH_pv = fitH/factorial(np.arange(fitOrder+1))
fitH_pv = fitH_pv[::-1]

fitV_pv = fitV/factorial(np.arange(fitOrder+1))
fitV_pv = fitV_pv[::-1]

labelH = r'$Q={:.1f}$'.format(fitH[0]) + '\n' +  \
         r"$Q'={:.2f}$".format(fitH[1]) + '\n' +   \
         r"$Q''={:.4f}$".format(fitH[2])

labelV = r'$Q={:.1f}$'.format(fitV[0]) + '\n' +  \
         r"$Q'={:.2f}$".format(fitV[1]) + '\n' +   \
         r"$Q''={:.4f}$".format(fitV[2])

fig, ax = plt.subplots(2,1)
ax[0].plot(dpa, np.polyval(fitH_pv, dpa), label=labelH)
ax[0].plot(dpa, qz[:,0], linestyle='None', marker='x')
ax[1].plot(dpa, np.polyval(fitV_pv, dpa), label=labelV)
ax[1].plot(dpa, qz[:,1], linestyle='None', marker='x')
ax[0].legend()
ax[1].legend()

plt.show()