__author__ = 'Sergey'

from scipy.integrate import simpson

from ocelot.cpbd.optics import trace_z, twiss
from ocelot.cpbd.beam import *
from ocelot.cpbd.elements import *
from ocelot.rad.undulator_params import *
from ocelot.common.ocelog import *
from ocelot.cpbd.transformations.transformation import TMTypes

_logger = logging.getLogger(__name__)


def I2_ID(L, h0):
    return L/2.*h0*h0


def I3_ID(L, h0):
    return 4.*L/(3*pi)*h0**3


def I4_ID(L, h0, lu):
    return -3.*L*lu*lu*h0**4/(32*pi)


def I5_ID(L, h0, lu, beta_xc, Dx0, Dxp0):
    # it is the same as I5_exact and I5_exact2
    # beta_c - beta_x at the center of ID
    # Dx0, Dxp0 - Dx and Dxp at the beginning of ID
    nn = int(L/lu)
    I = ((h0**3 * L)/(108000 * pi**5 * beta_xc) * (144000 * pi**4 * (Dx0**2 + beta_xc**2 * Dxp0**2) +
                                                   13500 * (-1)**nn * h0 * pi**3 * Dx0 * lu**2 + 15656 * h0**2 * lu**4 +
                                                   15 * (-76 + 225 * (-1)**nn) * h0**2 * pi * lu**4 +
                                                   150 * h0 * pi**2 * lu**2 * (480 * Dx0 + h0 * (4 * L**2 + 48 * beta_xc**2 - lu**2))))
    return I


def radiation_integrals(lattice, twiss_0, nsuperperiod=1):
    # TODO: add I4 for rectangular magnets I4 = Integrate(2 Dx(z)*k(z)*h(z), Z)

    n_points_element = 20

    tws_elem = twiss_0
    (I1, I2, I3, I4, I5) = (0., 0., 0., 0., 0.)
    h = 0.
    for elem in lattice.sequence:
        for tm in elem.first_order_tms:
            # TODO: Question: Why do we only have to split the Bend Main TM?
            if elem.__class__ in (SBend, RBend, Bend) and tm.tm_type == TMTypes.MAIN:
                Dx = []
                Hinvariant = []
                Z = []
                h = elem.angle/elem.l

                for z in np.linspace(0, elem.l, num=n_points_element, endpoint=True):
                    section_tms = elem.get_section_tms(delta_l=z, ignore_edges=True)
                    tws_z = tws_elem
                    for s_tm in section_tms:
                        tws_z = s_tm * tws_z
                    Dx.append(tws_z.Dx)
                    Z.append(z)
                    Hx = (tws_z.gamma_x*tws_z.Dx*tws_z.Dx + 2.*tws_z.alpha_x*tws_z.Dxp*tws_z.Dx
                          + tws_z.beta_x*tws_z.Dxp*tws_z.Dxp)
                    Hinvariant.append(Hx)
                #H = array(h)
                H2 = h*h
                H3 = np.abs(h*h*h)
                I1 += h*simpson(np.array(Dx), x=Z)
                I2 += H2*elem.l  # simpson(H2, x=Z)*nsuperperiod
                I3 += H3*elem.l  # simpson(H3, x=Z)*nsuperperiod
                I4 += h*(2*elem.k1 + H2)*simpson(np.array(Dx), x=Z)
                I5 += H3*simpson(np.array(Hinvariant), x=Z)
            tws_elem = tm * tws_elem
    # if abs(tws_elem.beta_x - twiss_0.beta_x)>1e-7 or abs(tws_elem.beta_y - twiss_0.beta_y)>1e-7:
    #    print( "WARNING! Results may be wrong! radiation_integral() -> beta functions are not matching. ")
        # return None
    return (I1*nsuperperiod, I2*nsuperperiod, I3*nsuperperiod, I4*nsuperperiod, I5*nsuperperiod)


class EbeamParams:
    def __init__(self, lattice, tws0, coupling=0.01, nsuperperiod=1):
        if tws0.__class__ is not Twiss:
            _logger.info(" EbeamParams: tws0 is not Twiss class.")
            return
        self.tws0 = tws0
        self.E = tws0.E

        if tws0.beta_x == 0 and tws0.beta_y == 0:
            tws = twiss(lattice, tws0)
            self.tws0 = tws[0]

        if self.E == 0.0:
            _logger.info("tws0.E is 0. Some parameters will be NaN")

        self.lat = lattice
        (I1, I2, I3, I4, I5) = radiation_integrals(lattice, self.tws0, nsuperperiod)
        self.I1 = I1
        self.I2 = I2
        self.I3 = I3
        self.I4 = I4
        self.I5 = I5
        self.Je = 2.0 + I4/I2
        self.Jx = 1.0 - I4/I2
        self.Jy = 1.0
        self.gamma = self.E/m_e_GeV
        self.sigma_e = self.gamma*np.sqrt(Cq * self.I3/(self.Je*I2))
        self.emittance = Cq*self.gamma*self.gamma * self.I5/(self.Jx * self.I2)
        self.U0 = Cgamma*(self.E*1000.0)**4*self.I2/(2.0*np.pi)

        self.Tperiod = nsuperperiod*lattice.totalLen/speed_of_light
        self.Length = nsuperperiod*lattice.totalLen
        self.tau0 = 2.0*self.E*1000.0*self.Tperiod/self.U0 if self.U0 != 0.0 else np.nan
        self.tau_e = self.tau0/self.Je
        self.tau_x = self.tau0/self.Jx
        self.tau_y = self.tau0/self.Jy
        self.alpha = self.I1/(speed_of_light*self.Tperiod)
        self.coupl = coupling
        self.emitt_x = self.emittance/(1.0 + self.coupl)
        self.emitt_y = self.emittance*self.coupl/(1.0 + self.coupl)
        self.sigma_x = np.sqrt((self.sigma_e*self.tws0.Dx)**2 + self.emitt_x*self.tws0.beta_x)
        self.sigma_y = np.sqrt((self.sigma_e*self.tws0.Dy)**2 + self.emitt_y*self.tws0.beta_y)
        self.sigma_xp = np.sqrt((self.sigma_e*self.tws0.Dxp)**2 + self.emitt_x*self.tws0.gamma_x)
        self.sigma_yp = np.sqrt((self.sigma_e*self.tws0.Dyp)**2 + self.emitt_y*self.tws0.gamma_y)

    def integrals_id(self):
        L = 0.
        self.I2_IDs = 0.
        self.I3_IDs = 0.
        self.I4_IDs = 0.
        self.I5_IDs = 0.
        for elem in self.lat.sequence:
            if elem.__class__ == Undulator:
                B = K2field(elem.Kx, lu=elem.lperiod)
                h0 = B*speed_of_light/self.E*1e-9
                tws = trace_z(self.lat, self.tws0, [L, L + elem.l/2.])
                i2 = I2_ID(elem.l, h0)
                i3 = I3_ID(elem.l, h0)
                i4 = I4_ID(elem.l, h0, elem.lperiod)
                i5 = I5_ID(elem.l, h0, elem.lperiod, tws[1].beta_x, tws[0].Dx, tws[0].Dxp)
                self.I2_IDs += i2
                self.I3_IDs += i3
                self.I4_IDs += i4
                self.I5_IDs += i5
            L += elem.l
        self.emit_ID = self.emittance * (1.+self.I5_IDs/self.I5)/(1+(self.I2_IDs - self.I4_IDs)/(self.I2 - self.I4))
        self.sigma_e_ID = self.sigma_e * np.sqrt((1. + self.I3_IDs / self.I3)/(1 + (2*self.I2_IDs + self.I4_IDs)/(2.*self.I2 + self.I4)))
        self.U0_ID = Cgamma*(self.E*1000)**4.*self.I2_IDs/(2.*np.pi)
        print("emittance with IDs = ", self.emit_ID*1e9, " nm*rad")
        print("sigma_e with IDs =   ", self.sigma_e_ID)
        print("U0 from IDs =        ", self.U0_ID,  "  MeV")

    def __str__(self):
        val = ""
        val += ("I1 =        " + str(self.I1))
        val += ("\nI2 =        " + str(self.I2))
        val += ("\nI3 =        " + str(self.I3))
        val += ("\nI4 =        " + str(self.I4))
        val += ("\nI5 =        " + str(self.I5))
        val += ("\nJe =        " + str(self.Je))
        val += ("\nJx =        " + str(self.Jx))
        val += ("\nJy =        " + str(self.Jy))
        val += ("\nenergy =    " + str(self.E) + "GeV")
        val += ("\ngamma =     " + str(self.gamma))
        val += ("\nsigma_e =   " + str(self.sigma_e))
        val += ("\nemittance = " + str(self.emittance*1e9) + " nm*rad")
        val += ("\nLength =    " + str(self.Length) + " m")
        val += ("\nU0 =        " + str(self.U0) + "  MeV")
        val += ("\nTperiod =   " + str(self.Tperiod*1e9) + " nsec")
        val += ("\nalpha =     " + str(self.alpha))
        val += ("\ntau0 =      " + str(self.tau0*1e3) + " msec")
        val += ("\ntau_e =     " + str(self.tau_e*1e3) + " msec")
        val += ("\ntau_x =     " + str(self.tau_x*1e3) + " msec")
        val += ("\ntau_y =     " + str(self.tau_y*1e3) + " msec")
        val += ("\nbeta_x =    " + str(self.tws0.beta_x) + " m")
        val += ("\nbeta_y =    " + str(self.tws0.beta_y) + " m")
        val += ("\nalpha_x =   " + str(self.tws0.alpha_x))
        val += ("\nalpha_y =   " + str(self.tws0.alpha_y))
        val += ("\nDx =        " + str(self.tws0.Dx) + " m")
        val += ("\nDy =        " + str(self.tws0.Dy) + " m")
        val += ("\nsigma_x =   " + str(self.sigma_x*1e6) + " um")
        val += ("\nsigma_y =   " + str(self.sigma_y*1e6) + " um")
        val += ("\nsigma_x' =  " + str(self.sigma_xp*1e6) + " urad")
        val += ("\nsigma_y' =  " + str(self.sigma_yp*1e6) + " urad\n")
        return val
