import numpy as np
import at
from matplotlib import pyplot as plt

print(f'at.__version__={at.__version__}')
print(f'np.__version__={np.__version__}')

tune_diff = np.zeros((1000,3))
for i in range(1,1001):
    the_length = i*0.001
    QF = at.Multipole("QF", 0.5, poly_b=np.array([0,1.2]), poly_a=np.zeros(2))
    
    drift_length = 0.5
    Dr = at.Drift("Dr", drift_length)
    HalfDr = at.Drift("Dr2", drift_length/2)
    Dr3 = at.Drift("Dr3", 0.075)
    Dr4 = at.Drift("Dr4", the_length)
    und = at.Wiggler('und', length=Dr4.Length, wiggle_period=Dr4.Length/99, b_max=0,Nmeth=2,Nstep=10,energy=1e9)
    
    QD = at.Quadrupole("QD", 0.5, -1.2)
    Bend = at.Dipole("Bend", 1, 2 * np.pi / 20)
    
    FODOcell = at.Lattice(
        [HalfDr, Bend, Dr, QF, Dr, Bend, Dr, QD, Dr3, Dr4],
        name="Simple FODO cell",
        energy=1e9,
    )
    
#    print(FODOcell.radiation_parameters().tunes)
#    print(FODOcell.circumference)
    
    uFODOcell = FODOcell.deepcopy()
    uFODOcell[-1] = und.deepcopy()
#    
#    print(uFODOcell.radiation_parameters().tunes)
#    print(uFODOcell.circumference)
    #print(uFODOcell.radiation_parameters().tunes-FODOcell.radiation_parameters().tunes)
    tune_diff[i-1,:] = uFODOcell.radiation_parameters().tunes-FODOcell.radiation_parameters().tunes


plt.plot(0.001*np.arange(1,1001),tune_diff)
plt.xlabel('Length [m]')
plt.ylabel('tune diff')
plt.grid()
plt.legend(["X", "Y"])
plt.show()