import numpy as np
import matplotlib.pyplot as plt

def get_random_walkers(Nwalkers, Ntimes):
	"""Generate Nwalkers for Ntimes.

	Args:
		Nwalkers (integer): No. of walkers
		Ntimes (integer): No. of times

	Returns:
		ndarray(Nwalkers,Ntimes+1): Different realizations of the walkers
	"""
	#
    # Steps
	walkers = np.random.choice( (-1,1), size=(Nwalkers,Ntimes+1) )
	#walkers = np.random.random( size=(Nwalkers,Ntimes+1) )
	#walkers = np.sign( walkers - 0.5 )

    # generate the walkers
	walkers[:,0] = 0
	walkers = np.cumsum(walkers, axis=1)
    #
	return walkers
# end function





def main(Nwalkers, Ntimes):
	#
    # Generate the walkers
	Walkers = get_random_walkers(Nwalkers, Ntimes)

	# Time
	tid = np.arange(Ntimes+1)
    
	# Plotter noen av walkerene
	for t in range(5):
		plt.plot(tid,Walkers[t,:])
		plt.xlabel("tidssteg")
		plt.ylabel("posisjon")
		plt.show()
	# end for

	# Plotter the pdfene ved tidene 10,100,500,1000
	for t in (10,100,500,1000):
		plt.hist(Walkers[:,t],bins='auto',density=True)
		plt.xlabel("posisjon x")
		plt.ylabel("p(x)")
		plt.title(f"t={t}".format(t=t))
		plt.show()
	# end for


	# Get mean and std and plot them
	mean = np.mean(Walkers,axis=0)
	std  = np.std(Walkers,axis=0)

	plt.plot(tid, mean, label='mean')
	plt.plot(tid, std,  label='std')
	plt.plot(tid, np.sqrt(tid),  label='sqrt(t)')
	plt.legend()
	plt.xlabel('tid')
	plt.show()
	#
# End function	

    









if __name__ == "__main__":

	NTIMES   = 1000
	NWALKERS = 1000

	main(NWALKERS,NTIMES)

# end if