#!/usr/bin/python3

### koda za generiranje rezultatov za seminar 

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
rc('font',**{'family':'serif','serif':['Computer Modern']})
rc('text', usetex=True)
from boris import *

Rz = 6378137 # radii zemlje [m]

q = e 
m = m_pr

K = 1e7 # kineticna energija 10MeV
K = K*e # v Jouleih

get_v0 = lambda K: c*np.sqrt(1-1/(K/(m*c**2)+1)**2)


x0 = np.array([2*Rz,0.,0.])
vpad = 45
v0 = get_v0(K)*np.array([0., np.sin(vpad*np.pi/180), np.cos(vpad*np.pi/180)])

E = lambda x: np.array([0.,0.,0.]) # elektricno polje je nula
dt = 1e-3
tsim = 20

X, V = boris(x0, v0, E, B_zemlja, dt, tsim, q=q,m=m)
x0 = np.array([4*Rz,0.,0.])
X2, V2 = boris(x0, v0, E, B_zemlja, dt, tsim, q=q,m=m)

def plot(X,X2,r=[2,4], fname=None):
    fig = plt.figure()
    ax = fig.subplots()
    ax.plot(X[:,0],X[:,1],label='${}R_z$'.format(r[0]))
    ax.plot(X2[:,0],X2[:,1],label='${}R_z$'.format(r[1]))
    ax.grid()
    ax.set_xlabel(r'$x$ [$R_z$]')
    ax.set_ylabel(r'$y$ [$R_z$]')
    ax.legend()
    if fname:
        plt.savefig(fname,bbox_inches='tight')
    plt.show()

plot3(X2/Rz, enote="$R_z$",zemlja=True)

plot(X/Rz,X2/Rz,fname='radii.pdf')