'''
renamed variables not to have a variable with the same name as the fuction the variable is used within
-- ME. 2018-07-23

added geometrical and normalized emittance calculations
-- ME. 2019-02-19
'''
from scipy import constants

# Proton mass in MeV/c
proton_mass_in_MeV=constants.physical_constants['proton mass energy equivalent in MeV'][0]
# Speed of light in m/s
c=constants.c

def beta(energy: float, mass=proton_mass_in_MeV):
    '''
    Parameters
    ----------
    energy:  float
             kinetic energy of the particle in MeV/c2
    mass:    float
             mass of the moving particle, optional, if not given mass of proton is used

    Returns
    -------
    relaivistic beta of the particle
    '''
    b=(1-gamma(energy, mass)**-2)**0.5
    return b   

def beta_c(energy, mass=proton_mass_in_MeV):
    '''
    Parameters
    ----------
    energy:  float
             kinetic energy of the particle in MeV/c2
    mass:    float
             mass of the moving particle, optional, if not given mass of proton is used

    Returns
    -------
    relaivistic c*beta of the particle
    '''    
    bc=c*(1-gamma(energy, mass)**-2)**0.5
    return bc   

def beta_from_gamma(gamma):
    '''
    Parameters
    ----------
    gamma:   float
             relaivistic gamma of the particle

    Returns
    -------
    relaivistic beta of the particle

    raises
    ------
    if given gamma is smaller than one raises a warning 
    '''    
    if gamma < 1:
        print('gamma cannot be less than one')
        beta = 0
    else:
        beta=(1-gamma**-2)**0.5
    return beta  

def gamma(energy, mass=proton_mass_in_MeV):
    '''
    Parameters
    ----------
    energy:  float
             kinetic energy of the particle in MeV/c2
    mass:    float
             mass of the moving particle, optional, if not given mass of proton is used

    Returns
    -------
    relaivistic gamma of the particle
    '''
    #if mass>0:
    g=1+(energy/mass)
    return g

def gamma_from_beta(beta):
    '''
    Parameters
    ----------
    beta:    float
             relaivistic beta of the particle

    Returns
    -------
    relaivistic gamma of the particle
    '''
    if beta == 1:
        from math import inf as math_inf
        gamma = math_inf 
    else:
        gamma = (1-beta**2)**-0.5
    return gamma  

def energy(gamma_beta, mass=proton_mass_in_MeV):
    '''
    Parameters
    ----------
    gamma_beta:  float
                 gamma or beta of the particle; if gamma_beta >= 1 the value is gamma, if gamma_beta < 1 the value is beta

    mass:        float
                 mass of the moving particle, optional, if not given mass of proton is used

    Returns
    -------
    returns kinetic energy of the particle in MeV.

    
    example: 
    sprel.energy(2.55, 938) 
    > 1453.8999999999999
    if gamma_beta < 1 the value is beta
    example:
    sprel.energy(.55, 938)
    > 185.13182200743233
    '''
    if gamma_beta >= 1: #assuming the value given is gamma
        e = (gamma_beta - 1) * mass
    elif gamma_beta < 1: #adduming the value given is beta
        e = mass * (-1 + 1 / (1 - gamma_beta ** 2) ** 0.5)
    return e

def emit_geo(u_up_array):
    '''
    parameters
    ----------
    u_up_array: array like
        a 2D array containing that has u_up_array[:,:2] as coordinates in u and u'
        
    returns
    -------
    geometrical emittance of the distribution/array in units of pi.mm.mrad
    '''
    u_shift = u_up_array[:,0]      #  x
    up_shift = u_up_array[:,1]     #  x'
    u = u_shift - np.average(u_shift)
    up = up_shift - np.average(up_shift)
    u2 = np.square(u)        #  x2
    u2ave = np.average(u2)   # <x2> 
    print(u2ave)
    up2 = np.square(up)      #  x'2
    up2ave = np.average(up2) # <x'2>
    print(up2ave)
    uup = np.multiply(u, up)      #  x.x'  
    uupave = np.average(uup) # <x.x'>
    uupave2 = uupave**2      # <x.x'>2
    print(uupave2)
    emit2 = (u2ave * up2ave) - uupave2
    emit = np.sqrt(emit2)
    return emit

def emit_norm(u_up_array, gamma):
    '''
    parameters
    ----------
    u_up_array: array like
        a 2D array containing that has u_up_array[:,:2] as coordinates in u and u'
    gamma: float
        relativistic gamma of the particles in the beam
    returns
    -------
    normalized emittance of the distribution/array in units of pi.mm.mrad
    '''
    #from ess import SP_Relativity as spr
    bt = spr.beta_from_gamma(gamma)
    geo_emit = emit_geo(u_up_array)
    emit = bt * gamma * geo_emit
    return emit