# ~ -------------------------------------------------
# ~ Typical imports
# ~ -------------------------------------------------
import numpy as np
from numba import int64, float64    # import the types
import numba as nb
import matplotlib.pyplot as plt
import time as tm
# ~ -------------------------------------------------
# ~ In newer python versions jitclass is inside numba.experimental instead of just numba
# ~ so this script tries both options 
# ~ -------------------------------------------------

try:
	from numba.experimental import jitclass
	print("importing jitclass from numba.experimental")
except:
	from numba import jitclass
	print("importing jitclass from numba")


# ~ -------------------------------------------------
# ~ Declaration of the variables and types used in the compilation process
# ~ Since python is usually interpreted and dynamically allocates memory, 
# ~ which costs a lot of time, this is the crucial point of speed up
# ~ -------------------------------------------------

spec = [
	('r',  float64[:]),
	('trajectory', float64[:,:]),
	('dt',  float64),
	('Diffusion', float64),
	('noise_constant',  float64),    
	('Ndt',  int64),    
	
]

# ~ -------------------------------------------------
# ~ Example of a class which is compiled as a whole
# ~ -------------------------------------------------

@jitclass(spec)
class brownian_particle:
	def __init__(self,NDT):
		self.r = np.zeros(2, dtype=float64)
		
		self.dt = 0.00001
		self.Diffusion = 0.01
		self.noise_constant = np.sqrt(2.0*self.dt*self.Diffusion)
		self.Ndt = NDT
		self.trajectory = np.zeros((self.Ndt,2), dtype=float64)
		#print(self.trajectory)
	def perform_noise(self):
		#print(f"action = {self.chosen_action}; dphi = {self.dphi_actions}")
		self.r[0] += self.noise_constant*np.random.normal()
		self.r[1] += self.noise_constant*np.random.normal()	
	def simulate(self):	
		for i in range(self.Ndt):
			self.perform_noise()
			self.trajectory[i,0] = self.r[0]
			self.trajectory[i,1] = self.r[1]
			# ~ self.trajectory[i] = self.r
			
			

		return self.trajectory
# ~ -------------------------------------------------
# ~ Example of a function outside of a class which is compiled as such
# ~ -------------------------------------------------		
@nb.njit
def someRandomFunction(x):
	return x*x

# ~ -------------------------------------------------
# ~ Main routine. Can in principle also be defined as njit function in- our outside class
# ~ for further speedup
# ~ -------------------------------------------------		

Nsteps = 100000

BP =  brownian_particle(Nsteps)	




T0 = tm.time()
#compiling and running once
BP.simulate()

T1 = tm.time()	
Tnumbafirst = T1-T0
print(f"total time of compilation and first run with numba= {Tnumbafirst} s")

BP.r[0] = 0.0
BP.r[1] = 0.0

T0 = tm.time()
trajectory = BP.simulate()
T1 = tm.time()	

Tnumba = T1-T0

trajectory = np.array(trajectory)

print(f"total time with numba without compilation = {Tnumba} s")

fig,ax = plt.subplots()
#ax.plot(trajectory[:,0],trajectory[:,1])
ax.set_aspect("equal")	


# ~ -------------------------------------------------
# ~ Below the same without numba for comparison
# ~ -------------------------------------------------		
	
class brownian_particle1:
	def __init__(self,NDT):
		self.r = np.zeros(2)
		self.dt = 0.00001
		self.Diffusion = 0.01
		self.noise_constant = np.sqrt(2.0*self.dt*self.Diffusion)
		self.Ndt = NDT

	def perform_noise(self):
		#print(f"action = {self.chosen_action}; dphi = {self.dphi_actions}")
		self.r[0] += self.noise_constant*np.random.normal()
		self.r[1] += self.noise_constant*np.random.normal()	

BP = brownian_particle1(Nsteps)
T0 = tm.time()
trajectory = []
for i in range(BP.Ndt):
	BP.perform_noise()
	trajectory.append([BP.r[0],BP.r[1]])
	
T1 = tm.time()	
Tnonumba = T1-T0
trajectory = np.array(trajectory)

print(f"total time without numba= {Tnonumba} s")

fig,ax = plt.subplots()
ax.plot(trajectory[:,0],trajectory[:,1])
ax.set_aspect("equal")