# Calculates optimum fraction of beam as a function of area of ellipse
# Please not that this algorithm is not fool proof as optimizer might
# get stuck due to discrete number of particles.

import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt

# Read in particles and center
x,xp,I = np.loadtxt('particles.txt',unpack=True)
Itotal = np.sum(I)
xave = np.sum(I*x)/Itotal
xpave = np.sum(I*xp)/Itotal
x = x-xave
xp = xp-xpave

# Calculate rms values
x2 = np.sum(I*x**2)/Itotal
xp2 = np.sum(I*xp**2)/Itotal
xxp = np.sum(I*x*xp)/Itotal
epsilon = np.sqrt(x2*xp2 - xxp**2)

# Print
print('epsilon_rms =',1e6*epsilon,'mm mrad')
print('I =',1e3*Itotal,'mA')

# Twiss parameters from rms
alpha = -xxp/epsilon
beta = x2/epsilon
gamma = xp2/epsilon

# Returns if particle at (x,xp) is inside ellipse defined by alpha,beta,epsilon.
def inside(x,xp,alpha,beta,epsilon):
    gamma = (1+alpha**2)/beta
    if x <= -np.sqrt(epsilon*beta) or x >= np.sqrt(epsilon*beta):
        return False
    xpmin = (-2*alpha*x - np.sqrt((2*alpha*x)**2-4*beta*(gamma*x**2-epsilon)))/(2*beta)
    xpmax = (-2*alpha*x + np.sqrt((2*alpha*x)**2-4*beta*(gamma*x**2-epsilon)))/(2*beta)
    if xp <= xpmin or xp >= xpmax:
        return False
    return True

# Returns number of particles ellipse defined by alpha, beta,
# epsilon. Ellipse shape (alpha, beta) are optimized to minimize particles outside.
def outside_count(p,args):
    alpha = p[0]
    beta = p[1]
    x = args[0]
    xp = args[1]
    epsilon = args[2]

    gamma = (1+alpha**2)/beta
    mask = np.logical_and(x > -np.sqrt(epsilon*beta), x < np.sqrt(epsilon*beta))
    xpmin = (-2*alpha*x[mask] - np.sqrt((2*alpha*x[mask])**2-4*beta*(gamma*x[mask]**2-epsilon)))/(2*beta)
    xpmax = (-2*alpha*x[mask] + np.sqrt((2*alpha*x[mask])**2-4*beta*(gamma*x[mask]**2-epsilon)))/(2*beta)
    mask2 = np.logical_and(xp[mask] > xpmin, xp[mask] < xpmax)
    return(len(xp)-len(xp[mask][mask2]))

frac_eps = []
frac_I = []
frac_Iopt = []
av = []
bv = []

# Calculate optimal ellipse shapes for 51 cases
N = 51
eps_rms = epsilon
alpha_rms = alpha
beta_rms = beta
gamma_rms = gamma
for i in range(N):
    frac_eps.append( 5*(N-i)/(N-1) )
    epsilon = frac_eps[-1]*eps_rms
    frac_I.append( 1-outside_count((alpha_rms,beta_rms),(x,xp,epsilon))/len(x) )
    
    x0 = [[alpha,beta],[1.1*alpha,beta],[alpha,1.1*beta]]
    res = minimize(outside_count,[alpha,beta],args=[x,xp,epsilon],method='Nelder-Mead',options={'initial_simplex':x0})
    if res.success == False:
        print('Error on round i =',i)
        print(res)
        exit(1)
    alpha = res.x[0]
    beta = res.x[1]
    gamma = (1+alpha**2)/beta
    av.append(alpha)
    bv.append(beta)
    frac_Iopt.append( 1-outside_count((alpha,beta),(x,xp,epsilon))/len(x) )

# Add zeros
frac_eps.append(0)
frac_I.append(0)
frac_Iopt.append(0)
av.append(0)
bv.append(0)

plt.plot(frac_eps,frac_I,'r',label='rms shape')
plt.plot(frac_eps,frac_Iopt,'b',label='optimized')
plt.legend()
plt.xlabel('Area of ellipse ($\\epsilon_{{rms}}$)')
plt.ylabel('Fraction of beam')
plt.savefig('epsfrac.pdf')

#----------------------------------------

# Plot particle distribution
plt.figure()
plt.plot(1e3*x,1e3*xp,'r,')

# Plot optimal ellipse for 1*Erms
i = 11
alpha = av[i]
beta = bv[i]
gamma = (1+alpha**2)/beta
epsilon = frac_eps[i]*eps_rms
phi = 0.5*np.arctan2( -2*alpha, beta - gamma )
H = 0.5*(beta+gamma)
rmajor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)+np.sqrt(H-1) )
rminor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)-np.sqrt(H-1) )

t = np.linspace(0,2*np.pi)
xell = rmajor*np.cos(t)*np.cos(phi) - rminor*np.sin(t)*np.sin(phi)
yell = rmajor*np.cos(t)*np.sin(phi) + rminor*np.sin(t)*np.cos(phi)
plt.plot(1e3*xell,1e3*yell,'k-')
print('')
print('Emittance =',frac_eps[i],'times E_rms,', 1e6*eps_rms*frac_eps[i], 'mm mrad')
print('Optimal current =',frac_Iopt[i],'times Itotal,',1e3*Itotal*frac_Iopt[i],'mA')
print('RMS current =',frac_I[i],'times Itotal,',1e3*Itotal*frac_I[i],'mA')

# Plot optimal ellipse for 4*Erms
i = 41
alpha = av[i]
beta = bv[i]
gamma = (1+alpha**2)/beta
epsilon = frac_eps[i]*eps_rms
phi = 0.5*np.arctan2( -2*alpha, beta - gamma )
H = 0.5*(beta+gamma)
rmajor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)+np.sqrt(H-1) )
rminor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)-np.sqrt(H-1) )

t = np.linspace(0,2*np.pi)
xell = rmajor*np.cos(t)*np.cos(phi) - rminor*np.sin(t)*np.sin(phi)
yell = rmajor*np.cos(t)*np.sin(phi) + rminor*np.sin(t)*np.cos(phi)
plt.plot(1e3*xell,1e3*yell,'k-')
print('')
print('Emittance =',frac_eps[i],'times E_rms,', 1e6*eps_rms*frac_eps[i], 'mm mrad')
print('Optimal current =',frac_Iopt[i],'times Itotal,',1e3*Itotal*frac_Iopt[i],'mA')
print('RMS current =',frac_I[i],'times Itotal,',1e3*Itotal*frac_I[i],'mA')

# Plot rms ellipse for 1*Erms
alpha = alpha_rms
beta = beta_rms
gamma = (1+alpha**2)/beta
epsilon = eps_rms
phi = 0.5*np.arctan2( -2*alpha, beta - gamma )
H = 0.5*(beta+gamma)
rmajor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)+np.sqrt(H-1) )
rminor = np.sqrt( 0.5*epsilon ) * ( np.sqrt(H+1)-np.sqrt(H-1) )

t = np.linspace(0,2*np.pi)
xell = rmajor*np.cos(t)*np.cos(phi) - rminor*np.sin(t)*np.sin(phi)
yell = rmajor*np.cos(t)*np.sin(phi) + rminor*np.sin(t)*np.cos(phi)
plt.plot(1e3*xell,1e3*yell,'k--')
rms_frac = 1-outside_count((alpha_rms,beta_rms),(x,xp,eps_rms))/len(x)

plt.xlabel('r (mm)')
plt.ylabel('r\' (mrad)')
plt.savefig('emitfit.pdf')

plt.show()