#!/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)

e = 1.602176565e-19 #osnovni naboj [C]
m_pr = 1.672621777e-27 #masa protona [kg]
m_el = 9.10938291e-31 #masa elektrona [kg]
c = 299792458 #hitrost svetlobe [m/s]

def boris(x0, v0, E, B, dt, tdur, 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 = int(tdur/dt)

    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

def B_zemlja(x):
    ''' model zemljinega magnentnega polja aproksimiran z 
    magnetnim dipolom '''
    B0 = 3.07e-5 # [T]
    Rz = 6378137 # radii zemlje [m]

    r = np.sqrt(np.sum(x**2))
    k = -B0*Rz**3/r**5
    return k*np.array([3*x[0]*x[2],3*x[1]*x[2],2*x[2]**2-x[0]**2-x[1]**2])

def Wk(m,V):
    return 0.5*m*np.sum(V**2,1)

def plot3(X,enote='m',zemlja=False):
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    if zemlja:
        u = np.linspace(0, 2 * np.pi, 100)
        v = np.linspace(0, np.pi, 100)
        x = np.outer(np.cos(u), np.sin(v))
        y = np.outer(np.sin(u), np.sin(v))
        z = np.outer(np.ones(np.size(u)), np.cos(v))
        ax.plot_surface(x,y,z,zorder=-4)
        #ax.set_box_aspect(np.ptp(X,axis=0))

    ax.plot(X[:,0],X[:,1],X[:,2],zorder=4)
    ax.set_xlabel(r'$x$ [{}]'.format(enote))
    ax.set_ylabel(r'$y$ [{}]'.format(enote))
    ax.set_zlabel(r'$z$ [{}]'.format(enote))

    plt.show()

def rplot(X, enote='m', xy=[0,1],os='xy', fname=None):
    fig = plt.figure()
    ax = fig.subplots()
    ax.plot(X[:,xy[0]],X[:,xy[1]], 'k')
    ax.grid()
    dx = X[-1,xy[0]]-X[-2,xy[0]]
    dy = X[-1,xy[1]]-X[-2,xy[1]]
    #ax.arrow(X[-1,xy[0]],X[-1,xy[1]],dx,dy,head_width=0.1,head_length=0.2,
    #        overhang=0.3, zorder=10,color='k')
    ax.set_xlabel(r'${}$ [{}]'.format(os[0],enote))
    ax.set_ylabel(r'${}$ [{}]'.format(os[1],enote))
    if fname:
        plt.savefig(fname,bbox_inches='tight')
    plt.show()



if __name__ == "__main__": # testni del kode
    E = lambda x: np.array([0.0,0.2,0.3])
    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,20)

    #ax.quiver(X[0:-1:20,0],X[0:-1:20,1],X[0:-1:20,2],V[0:-1:20,0],V[0:-1:20,1],V[0:-1:20,2], color='black',normalize=True,length=.5)
    # quiver res naredi grdo sliko
    plot3(X)