# -*- coding: utf-8 -*-
"""
General elements

@author: Alexis Gamelin
@date: 15/01/2020
"""

import numpy as np
from scipy.constants import c, e

        
class Synchrotron:
    """ Synchrotron class to store main properties. """
    def __init__(self, h, L, E0, particle, **kwargs):
        self.particle = particle
        self._h = h
        self.L = L
        self.E0 = E0 # Nominal (total) energy of the ring [eV]
        
        self.ac = kwargs.get('ac') # Momentum compaction factor
        self.U0 = kwargs.get('U0') # Energy loss per turn [eV]
        self.tau = kwargs.get('tau') # X/Y/S damping times [s]
        self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread
        self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s]
        self.tune = kwargs.get('tune') # X/Y/S tunes
        self.emit = kwargs.get('emit') # X/Y emittances in [m.rad]
        self.chro = kwargs.get('chro') # X/Y (non-normalized) chromaticities
        self.mean_optics = kwargs.get('mean_optics') # Optics object with mean values
                
    @property
    def h(self):
        """Harmonic number"""
        return self._h
    
    @h.setter
    def h(self,value):
        self._h = value
        self.L = self.L  # call setter
        
    @property
    def L(self):
        """Ring circumference [m]"""
        return self._L
    
    @L.setter
    def L(self,value):
        self._L = value
        self._T0 = self.L/c
        self._f0 = 1/self.T0
        self._omega0 = 2*np.pi*self.f0
        self._f1 = self.h*self.f0
        self._omega1 = 2*np.pi*self.f1
        self._k1 = self.omega1/c
        
    @property
    def T0(self):
        """Revolution time [s]"""
        return self._T0
    
    @T0.setter
    def T0(self, value):
        self.L = c*value
        
    @property
    def f0(self):
        """Revolution frequency [Hz]"""
        return self._f0
    
    @f0.setter
    def f0(self,value):
        self.L = c/value
        
    @property
    def omega0(self):
        """Angular revolution frequency [Hz rad]"""
        return self._omega0
    
    @omega0.setter
    def omega0(self,value):
        self.L = 2*np.pi*c/value
        
    @property
    def f1(self):
        """Fundamental RF frequency [Hz]"""
        return self._f1
    
    @f1.setter
    def f1(self,value):
        self.L = self.h*c/value
        
    @property
    def omega1(self):
        """Fundamental RF angular frequency[Hz rad]"""
        return self._omega1
    
    @omega1.setter
    def omega1(self,value):
        self.L = 2*np.pi*self.h*c/value
        
    @property
    def k1(self):
        """Fundamental RF wave number [m**-1]"""
        return self._k1
    
    @k1.setter
    def k1(self,value):
        self.L = 2*np.pi*self.h/value
    
    @property
    def gamma(self):
        """Relativistic gamma"""
        return self._gamma

    @gamma.setter
    def gamma(self, value):
        self._gamma = value
        self._beta = np.sqrt(1 - self.gamma**-2)
        self._E0 = self.gamma*self.particle.mass*c**2/e

    @property
    def beta(self):
        """Relativistic beta"""
        return self._beta

    @beta.setter
    def beta(self, value):
        self.gamma = 1/np.sqrt(1-value**2)
        
    @property
    def E0(self):
        """Nominal (total) energy of the ring [eV]"""
        return self._E0
    
    @E0.setter
    def E0(self, value):
        self.gamma = value/(self.particle.mass*c**2/e)

    @property
    def eta(self):
        """Momentum compaction"""
        return self.ac - 1/(self.gamma**2)
    
    @property
    def sigma(self):
        """RMS beam size at equilibrium"""
        sigma = np.zeros((4,))
        sigma[0] = (self.emit[0]*self.mean_optics.beta[0] + self.mean_optics.disp[0]**2*self.sigma_delta)**0.5
        sigma[1] = (self.emit[0]*self.mean_optics.alpha[0] + self.mean_optics.dispp[0]**2*self.sigma_delta)**0.5
        sigma[2] = (self.emit[1]*self.mean_optics.beta[1] + self.mean_optics.disp[1]**2*self.sigma_delta)**0.5
        sigma[3] = (self.emit[1]*self.mean_optics.alpha[1] + self.mean_optics.dispp[1]**2*self.sigma_delta)**0.5
        return sigma