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

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


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


def main():
    vel = [0., 1., 0.]
    x   = [1., 0., 0.]
    t = [0.]
    dt = 1.e-1
    print('initial radius = ', sum(np.power(x, 2)))
    for i in np.arange(0, 2000):
        #t = [0.1*i]
        t = t[-1]
        t = np.append(t, dt*i)
        y = np.hstack([x, vel])
        #print('y: ', y)
        z = odeint(func, y, t, args=(q, m))
        x,vel  = np.hsplit(z[-1], 2)
        #print('new position: ', x_new)
        #print(i, np.linalg.norm(x))
        if i%100 == 0: print(i, 'at t = ', t[-1],\
                             ' new radius = ', \
                             np.linalg.norm(x), 'x = ', x)
    #print(t)

main()