#!/usr/bin/env python
"""Try out odeint for a second-order differential equation, namely the equation of motion:

a = (q/m)*(v cross B)

where 

B is constant and given. 

The idea is to have a way of testing the accuracy of the solver odeint. 

How can this be done? Try to implement it.

"""



import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

B_field = [0.,0.,1.]
q = -1.
m = 1.


def func(y,dt, q, m):
    x,v = np.hsplit(y, 2)
    xp = v
    vp = np.multiply(q/m, np.cross(v, B_field))
    z = np.hstack([xp, vp])
    return z


def main():
    vel = [0., 1., 0.]
    x   = [1., 0., 0.]
    t = np.linspace(0,100,2000)
    y = np.hstack([x, vel])
    z = odeint(func, y, t, args=(q, m))
    x  = z[:, 0:3]
    v  = z[:, 3:6]
    print(np.shape(x))
    print(np.shape(v))
    print(x[-1, :], v[-1, :])
    print(np.linalg.norm(x, axis=1))
    print(np.linalg.norm(v, axis=1))

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

main()