import math
import numpy
from scipy.interpolate import RectBivariateSpline
from scipy.integrate import odeint
from matplotlib.pyplot import *
            
class MvtChampBaxial:
    def __init__(self,nom):
        self.Atheta = numpy.loadtxt(nom+"-Atheta.txt")
        self.AthetaDr = numpy.loadtxt(nom+"-AthetaDr.txt")
        self.AthetaDz = numpy.loadtxt(nom+"-AthetaDz.txt") 
        self.Br = numpy.loadtxt(nom+"-Br.txt")
        self.Bz = numpy.loadtxt(nom+"-Bz.txt")
        self.r = numpy.loadtxt(nom+"-r.txt")
        self.z = numpy.loadtxt(nom+"-z.txt")
        self.f_A=RectBivariateSpline(self.z,self.r,self.Atheta.transpose())
        self.f_ADr=RectBivariateSpline(self.z,self.r,self.AthetaDr.transpose())
        self.f_ADz=RectBivariateSpline(self.z,self.r,self.AthetaDz.transpose())
        self.f_Bz=RectBivariateSpline(self.z,self.r,self.Bz.transpose())
        self.f_Br=RectBivariateSpline(self.z,self.r,self.Br.transpose())
            
    def lignesB(self,color,density):
        streamplot(self.z,self.r,self.Bz,self.Br,color=color,density=density)
        xlabel("z")
        ylabel("r")
            
    def trajectoire(self,alpha,r0,z0,pr0,pz0,tmax,te):
        ptheta = alpha*r0*self.f_A(z0,r0)
        def equation(y,t):
            r=y[0]
            z=y[1]
            A = self.f_A(z,r)
            if ptheta==0:
                u = -alpha*A
                dpr = u*alpha*self.f_ADr(z,r)
                dpz = u*alpha*self.f_ADz(z,r)
                dtheta = -alpha*(self.f_Bz(z,r)-self.f_ADr(z,r))
            else:
                u = ptheta/r-alpha*A
                dpr = u*(ptheta/(r*r)+alpha*self.f_ADr(z,r))
                dpz = u*alpha*self.f_ADz(z,r)
                dtheta = ptheta/(r*r)-alpha*(self.f_Bz(z,r)-self.f_ADr(z,r))
            return [y[2],y[3],dpr,dpz,dtheta]
        t=numpy.arange(0.0,tmax,te)
        y = odeint(equation,[r0,z0,pr0,pz0,0.0],t,rtol=1e-6,atol=1e-6)
        yt = y.transpose()
        rt = yt[0]
        zt = yt[1]
        thetat = yt[4]
        return [rt,zt,thetat]
            
    def trajectoire_newton(self,alpha,r0,z0,pr0,pz0,tmax,te):
        if r0==0.0:
            r0 = 1e-10
        def equation(y,t):
            r=y[0]
            z=y[1]
            theta=y[2]
            Br = self.f_Br(z,r)
            Bz = self.f_Bz(z,r)
            dpr = r*y[5]*y[5]+alpha*r*y[5]*Bz
            dptheta = 1/r*(-2*y[3]*y[5]+alpha*(y[4]*Br-y[3]*Bz))
            dpz = -alpha*r*y[5]*Br
            return [y[3],y[4],y[5],dpr,dpz,dptheta]
        t=numpy.arange(0.0,tmax,te)
        y = odeint(equation,[r0,z0,0.0,pr0,pz0,0.0],t,rtol=1e-6,atol=1e-6) 
        yt = y.transpose()
        rt = yt[0]
        zt = yt[1]
        thetat = yt[2]
        return [rt,zt,thetat]