# -*- coding: utf-8 -*-
"""
Created on Mon Nov  5 14:13:29 2018

@author: Ondrej Ficker, Pravesh Dhyani
"""

import numpy as np
import matplotlib.pyplot as plt

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

shot_no = 18970
identifier = "loop_voltage"
# create data cache in the 'golem_cache' folder
ds = np.DataSource('golem_cache')
#Create a path to data and download and open the file
base_url = "http://golem.fjfi.cvut.cz/utils/data/"
data_file = ds.open(base_url + str(shot_no) + '/' + identifier)
#Load data from the file and plot to screen and to disk
data_Ul = np.loadtxt(data_file)


identifier = "loop_voltage"
data_file = ds.open(base_url + str(shot_no) + '/' + identifier)
#Load data from the file and plot to screen and to disk
data_Ul = np.loadtxt(data_file)

identifier='electron_density';
data_file = ds.open(base_url + str(shot_no) + '/' + identifier)
data_ne = np.loadtxt(data_file)

identifier='plasma_current';
data_file = ds.open(base_url + str(shot_no) + '/' + identifier)
data_Ip = np.loadtxt(data_file)

identifier='hxr';
data_file = ds.open(base_url + str(shot_no) + '/' + identifier)
data_hxr = np.loadtxt(data_file)

#some constants
c=3e8; k=1.38e-23; eV=11400; e=1.6e-19; eps=8.8e-12; me=9.109e-31; Zeff=5.0
Te=20
sh=18970 #setting shot number
beg=11000 #setting beginning of time axis
end=27000 #setting the end of the time axis
lng=end-beg

#loading data
dens=data_ne
Ul=data_Ul
Ip=data_Ip
hxr=data_hxr


lnlam=16; vte=4.19e5*np.sqrt(Te)
print(dens.shape)
print(data_Ul.shape)
print(Ip.shape)
print(hxr.shape)

E=data_Ul[beg:end,1]/(3.14*0.4*2)

#n=data_Ul[beg:end,1]

 
vc=np.zeros((lng))
I=np.zeros((lng))
err=np.zeros((lng))
intdr=np.zeros((lng))
vte=4.19e5*np.sqrt(Te)

#print shape(E), shape(n), shape(dens)

#fast averaging routine
def movingaverage(values,window):
    weigths = np.repeat(1.0, window)/window
    smas = np.convolve(values, weigths, 'valid')
    return smas 
# setting time axis
t=data_Ul[beg:end,0]
#interpolating density to general 1 MHz time axis
f=interpolate.interp1d(dens[:,0], dens[:,1])
n=f(t)
g=interpolate.interp1d(hxr[:,0], hxr[:,1])
hx=g(t)

#critical density and Dreicer field
vc=np.sqrt((1+0.5*Zeff)*n*e**3*lnlam/((1.0/0.43)*eps**2*4*pi*me*E))
nu_e=2.91e-6*n*lnlam*(Te)**(-1.5)
ed=e**3*n*(1+Zeff/2.0)*lnlam/(4*3.14*eps**2*me*vte**2)

#KB rate
nu_dr=0.3*nu_e*np.exp(np.log(E/ed)*(-3/16*(Zeff+1)))*np.exp(-1/(4*E/ed)-(np.sqrt((Zeff+1)/(E/ed))))
#print nu_dr, shape(nu_dr)

intdr=1e-6*np.cumsum(np.nan_to_num(nu_dr), dtype=float) #integrating Dreicer rate - overall RE number

Ec=e**3*n/(eps**2*4*pi*me*c**2)
#print max(Ec[:])
for i in range(0,len(vc[:])):  #Integration from critical velocity to c - rough estimation of immediate n_RE
    I[i],err[i]=quad(lambda x:n[i]*np.sqrt(1/np.pi/vte**2)*exp(-x**2/vte**2),vc[i],c)
fig, ax = plt.subplots(figsize=(12,20),nrows=4, sharex=True)

#some smoothing 
I=movingaverage(I,20)
hx=movingaverage(hx,20)
Ul=movingaverage(Ul[:,1],20)
nu_dr=movingaverage(nu_dr,20)
t=1000*t
#print max(intdr), intdr
#ploting of important parameters
ax[0].set_xlim(11,29)
ax[0].plot(t[:len(Ul)],100*Ul[beg:end], color='0.3', lw=2, label=r'$100*U_{loop}$')
ax[0].plot(t,Ip[beg:end,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(t,n, color=(0.1,0.9,0.4), lw=2, label=r'$n_{e}$')
ax[1].legend()
   
ax[1].set_ylabel(r"$n_{e}\,[\mathrm{m^{-3}}]$")
ax[2].semilogy(t[:len(I)-1000],I[:-1000], color=(0.8,0.1,0.8), lw=2, label=r'$n_{RE}$')
ax[2].semilogy(t[:len(nu_dr)-1000],nu_dr[:-1000], color=(0.1,0.1,0.8), lw=2, label=r'$\gamma_{KB}$')
ax[2].semilogy(t[:len(intdr)-1000],intdr[:-1000], color=(0.8,0.7,0.8), lw=2, label=r'$n_{RE,dr}$')
ax[2].set_ylim(1e-3,1e17)
ax[2].legend()
ax[2].set_ylabel(r"$\gamma_{KB}\,[\mathrm{m^{-3}s^{-1}}], n_{RE}\,[\mathrm{m^{-3}}]$")

ax[3].plot(t[:len(I)],hx,'r',label='HXR')
ax[3].legend()
ax[3].set_ylabel("Intensity [a.u.]")
ax[3].set_xlabel(r"$t\,\mathrm{[ms]}$")
    
plt.show()
plt.savefig("Re_discharge{}.png".format(sh))
#print I