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

Q = 1.602176565e-19

def boris(x0, v0, E, B, dt, q=1, m=1):
    ''' Borisov algoritem zabjega skoka za
    integracijo diferencialne enačbe

    x0 = [x,y,z] : vektor zacetne pozicije
    v0 = [vx,vy,vz] : vektor zacetne hitrosti
    E : funkcija, ki daje jakost polja E(x)
    B : funkcija, ki daje jakost polja B(x)
    dt: casovni korak '''

    duration = 1000

    X = np.zeros((duration,3))
    V = np.zeros((duration,3))
    x = x0
    v = v0

    q_prime = dt * q * 0.5 / m
    
    # https://en.wikipedia.org/wiki/Particle-in-cell
    for time in range(duration):
        h = q_prime*B(x)
        s = 2*h / (1+np.dot(h,h))
        u = v + q_prime*E(x)
        u_ = u + np.cross(u+np.cross(u,h),s)
        
        v = u_ + q_prime * E(x)
        x = x + dt*v
        V[time,:] = v
        X[time,:] = x

    return X, V

E = lambda x: np.array([0.0,0.0,0.])
B = lambda x: np.array([0.,0.0,1.*x[0]])

x0 = np.array([-1.,0.,0.])
v0 = np.array([0.,1.,0.])

dt = 1e-2

X, V = boris(x0,v0,E,B,dt)

ax = plt.figure().add_subplot(projection='3d')
ax.plot(X[:,0],X[:,1],X[:,2])
plt.xlabel(r'$x$ label')
plt.show()