import numpy as np
import matplotlib.pyplot as plt

class StochastikIterator(object):
	def __init__(self, r, num_simulations, chi0, std_deviation, maxiter=None):
		self._r = r
		self._num_simulations = num_simulations
		self._chi0 = chi0
		self._std_deviation = std_deviation
		self._maxiter=maxiter
		self._chi = None
		self._currentiter = None

	def __iter__(self):
		self._currentiter = 0
		self._chi = np.array([self._chi0] * self._num_simulations, dtype=np.float64)
		return self

	def __next__(self):
		if(self._maxiter and self._currentiter == self._maxiter):
			raise StopIteration()
		self._currentiter += 1


		result = self._chi.copy()
		self._chi *= np.exp(r - np.random.normal(scale=self._std_deviation, size=self._num_simulations))

		return result



	def do(self, iters=None):
		if(not (iters or self._maxiter)):
			raise ValueError("iters or maxiter must be specified")

		if(iters):
			until = iters
		if(self._maxiter and (not iters or (iters and self._maxiter < iters))):
				until = self._maxiter


		self._currentiter = 0
		self._chi = np.array([self._chi0] * self._num_simulations, dtype=np.float64)
		for i in range(until):
			self._chi *= np.exp(r - np.random.normal(scale=self._std_deviation, size=self._num_simulations))

		return self._chi.copy()



mean_chi = lambda r: np.mean(StochastikIterator(r, 500, 100, 0.25).do(70))

for r in [0.05, 0, -0.05]:
	print("r =", r, "\t<chi> = ", mean_chi(r))


R = np.arange(-4, 4, 0.01)

M = [mean_chi(r) for r in R]

plt.plot(R, M, label="$<\\chi>$")
plt.show()