#!/usr/bin/env python

from numpy import *
from math import *
import Gnuplot, Gnuplot.funcutils
import time
from scipy.integrate import odeint
from numpy.linalg import norm

B_field = [0.,0.,1.]


class particle(object):
    
    def __init__(self,pos):
	self.trajectory = []
        self.move(pos)
	#self.position = pos
	#self.trajectory = pos

    def __str__(self):
	rep = "particle is at: " + repr(self.position) + "\n"
	rep2 = "its trajectory was" + repr(self.trajectory)
	return rep + rep2

    def move(self,pos):
	self.position = pos
	self.trajectory.append(pos)


    def plot_trajectory(self):
	"""Plot the trajectory of the particle using gnuplot
        """
    	g = Gnuplot.Gnuplot()
	#x,y,z = array_split(self.trajectory,3,axis=0)
	#print self.trajectory, x,y,z
	#g.splot(x,y,z)
        g.reset()
        #string = 'set title "time: ' + str(t) +' s"'
        #g(string)
        g('set ylabel "y-axis [arb. units]"')   #this is how you access gnupot commands
        g('set xlabel "x-axis [arb. units]"')
        g('set zlabel "z-axis [arb. units]"')
        data_rho = Gnuplot.Data(self.trajectory, with_ ='line', title = 'rho')
	g.splot(data_rho)
	#z = time.sleep(20.00) 	#wait a little so you can look at it

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



def main():
    x = [1.,0.,0.]
    p = particle(x)
    vel = [0.,1.,0.1]
    t = [0.]
    dt = 1.e-1
    print 'initial radius = ', sum(power(x,2))
    #t = [0.1,0.2,0.4,0.6,0.8,1.]
    for i in xrange(0,2000):
        #t = [0.1*i]
        t = t[-1]
        t = append(t,dt*i)
        y = hstack([x,vel])
        #print 'y: ', y
        z = odeint(func,y,t)
        x,vel  = hsplit(z[-1],2)
        p.move(x)
        #print 'new position: ', x_new
        #print i, norm(x)
        if i%100 == 0: print i, 'at t = ', t[-1],' new radius = ', norm(x), 'x = ', x
    #print t
    p.plot_trajectory()
    raw_input('Please press the return key to continue...\n')

main()