import h5py
import numpy as numpy
import matplotlib as mplib
import matplotlib.pyplot as plt
import sys, getopt
import math


# read commandline arguments, first
fullCmdArguments = sys.argv

# - further arguments
filename = fullCmdArguments[1]

f=h5py.File(filename,'r')
r=f['/data/part/R']
z=f['/data/part/Z']
thet=f['/data/part/THET']
ur=f['/data/part/UR']
uz=f['/data/part/UZ']
uthet=f['/data/part/UTHET']
gamma=numpy.sqrt(1+ur[()]**2+uz[()]**2+uthet[()]**2)

vlight=299792458
qe=1.60217662E-19
me=9.109383E-31
eps_0=8.85418781762E-12
kb=1.38064852E-23
input00=f['/data/input.00']
B0 =input00.attrs['B0']
Rcurv =input00.attrs['Rcurv']
n0 =input00.attrs['n0']
dt =input00.attrs['dt']
it0d =input00.attrs['it0d']
it2d =input00.attrs['it2d']
itparts = input00.attrs['itparts']
nplasma = input00.attrs['nplasma']
msim = input00.attrs['msim']
qsim = input00.attrs['qsim']
lz   = f['/data/input.00/lz']


omepe=math.sqrt(abs(n0)*qe**2/(me*eps_0))
omece=qe*B0/me

# Normalizations
tnorm=1/omepe
rnorm=vlight*tnorm
bnorm=B0
enorm=vlight*bnorm
phinorm=enorm*rnorm
vnorm=vlight

rgrid=f['/data/var1d/rgrid'][()]*rnorm
zgrid=f['/data/var1d/zgrid'][()]*rnorm

x=r[()]*rnorm*numpy.cos(thet[()])
y=r[()]*rnorm*numpy.sin(thet[()])
fig, ax = plt.subplots(1,2)

for i in range(132300,132310):
    ax[0].plot(z[0:800,i]*rnorm,r[0:800,i]*rnorm,label='part %d' %i)
    ax[1].plot(x[0:800,i],y[0:800,i],label='part %d' %i)

ax[0].legend()
ax[0].grid()
ax[0].set_xlabel('Z [m]')
ax[0].set_ylabel('R [m]')
ax[1].legend()
ax[1].grid()
ax[1].set_xlabel('X [m]')
ax[1].set_ylabel('Y [m]')
plt.show()

#plt.savefig(filename+'_trajectories.pdf',orientation='landscape',)
#print(gamma)