# -*- coding: utf-8 -*-
"""
Created on Mon Apr 20 12:18:20 2015

@author: Ondřej

Requirements: HXR,loopvol,plasma current and electron density for selected shot,
all data may be downloaded from http://golem.fjfi.cvut.cz/shots/[number of shot]/
"""

from math import *
import numpy as np
import matplotlib.pyplot as plt
#mport pylab as pyl
from scipy.integrate import quad
from scipy import interpolate
#rom matplotlib.colors import LogNorm

#constants
c=3e8; k=1.38e-23; eV=11400; e=1.6e-19; eps=8.8e-12; me=9.109e-31; Zeff=3.0
Te=30
sh=18970 #setting shot number
BD=12000 #setting beggining of data to be used for the calculation of acc. (BD)
#end=30000 #setting end of data to be used for the calculation of acc.
Ul=loadtxt('loop_voltage_{}.txt'.format(sh))
Ip=loadtxt('plasma_current_{}.txt'.format(sh))
hxr=loadtxt('hxr_{}.txt'.format(sh))


def Borispusher(V,X,E,B,Q,m,dt):
    gamma=1/np.sqrt(1-np.multiply(V,V)/c**2)
    El=Q*dt/(m*2)*E;Bl=Q*dt/(m*2*gamma)*B
    u=gamma*V
    u=u+El
    uu=u+2*np.cross((u+np.cross(u,Bl)),Bl)/(1+Bl*Bl)
    unew=uu+El
    Vnew=unew/np.sqrt(1+unew**2/c**2)
    Xnew=X+Vnew*dt
    #xnew2=X+unew*dt
    return Vnew, Xnew
    
#Ul=np.loadtxt('loop_voltage_14580.txt')
#plt.plot(Ul[:,0],Ul[:,1])

#setting time steps and fields - for general purpose, field from Uloop further
dt=1e-6
T=28e-3
N=T/dt
E=np.zeros(3); B=np.zeros(3); T=np.zeros(N); X=np.zeros((3,N)); V=np.zeros((3,N)); R=np.zeros(N); phi=np.zeros(N)
E[0]=0;E[1]=0;E[2]=0
B[0]=0.0;B[1]=0.0;B[2]=0.0
V[0]=0;V[1]=0;V[2]=0
X[0]=0;X[1]=0;X[2]=0
#Ba=10;Ea=0.0
Q=-1.6e-19
m=9.109e-31
c=3e8  

# run Boris-Bunemann particle integrator for electric field due to Uloop 
for i in range(0,int(N)-1):
    E[0]=Ul[i+BD,1]/(2*np.pi*0.4)
    print E[0]
    V[:,i+1],X[:,i+1]=Borispusher(V[:,i],X[:,i],E,B,Q,m,dt)
    T[i]=i*dt
    

#fast averaging routine
def movingaverage(interval, window_size):
    window= np.ones(int(window_size))/float(window_size)
    return np.convolve(interval, window, 'same')
t=Ul[:,0]*1000 #setting millisecond time axis

#interpolating density to general 1 MHz time axis
hxr[:,1]=movingaverage(hxr[:,1],20)
Ul[:,1]=movingaverage(Ul[:,1],20)

#Plotting of important parameters and energy of accelerated particle

fig, ax =plt.subplots(nrows=3,sharex=True)
ax[0].set_xlim(10,35)
ax[0].plot(t[:],100*Ul[:,1], color='0.3', lw=2, label=r'$100*U_{loop}$')
ax[0].plot(t,Ip[:,1], color=(254/255.0,80/255.0,0.0), lw=2, label=r'$I_{p}$')
ax[0].set_ylim(0,1400)
ax[0].legend()
ax[0].set_ylabel(r'$I_{p} \, \mathrm{[A]}$, $U_{loop}\, \mathrm{[0.01*V]}$')

ax[1].plot(hxr[:-1,0]*1000,hxr[:-1,1],'r',label='HXR')
ax[1].legend()
ax[1].set_ylabel("Intensity [a.u.]")
ax[1].set_xlabel(r"$t\,\mathrm{[ms]}$")
print shape(t[BD:]), shape(V[1,:])
ax[2].plot(t[BD:],0.001*c**2*m*(1/np.sqrt(1-(V[0,:]**2+V[1,:]**2+V[2,:]**2)/c**2)-1)/-Q,'g',label='Vacuum acc.')
ax[2].legend()
ax[2].set_ylabel(r"$E$ [keV]")
ax[2].set_xlabel(r"$t\,\mathrm{[ms]}$")

#displaying and saving the figure
plt.show()
plt.savefig("Re_discharge_acc{}.png".format(sh))
#print I