import numpy as np
import matplotlib.pyplot as plt

# Question 2:
def MC_step(n, E, T):
    N = np.shape(n)[0]
    i = np.random.randint(N)
    j = np.random.randint(3)
    shift = np.random.randint(2)
    if shift == 0:
        shift = -1 # shift = - 1 implies that we propose n -> n-1, shift = + 1 implies that we propose n -> n+1.
    if n[i,j] + shift > 0:
        delta_E = 0.5 * np.pi ** 2 * (2 * shift * n[i,j] + 1)
        if delta_E <= 0: # Metropolis
            n[i,j] += shift
            E += delta_E
        elif delta_E > 0 and T > 0:
            r = np.random.random() # Metropolis
            if r < np.exp(- delta_E / T):
                n[i,j] += shift
                E += delta_E
    return E
    
# Question 3: approximately 25000 steps are required to reach the stationary state for N=1000: M_{b-i}=25000.
# M_{b-i} scales linearly with N approximately, because the atoms do not interact.
N = 1000
M = 200000
T = 10.
n = np.ones((N, 3), dtype=int)
E = np.zeros(M + 1)
e = 0.5 * np.pi ** 2 * np.sum(n ** 2) # current energy
E[0] = e
for k in range(M):
   e = MC_step(n, e, T)
   E[k+1] = e
   
plt.figure()
plt.plot(E)
plt.xlabel('Number of steps')
plt.ylabel('Energy')
plt.grid()
plt.show()
            
# Question 4: we observe the zero point energy when T->0 and we recover the equipartition theorem when T->+infinity.
M = 250000 # ten times the burn-in phase
N = 1000
Mbi = 25000 # Burn-in phase
Tarr = np.arange(0, 22., 2.)
Emean = np.zeros_like(Tarr)
for t, T in enumerate(Tarr):
    n = np.ones((N, 3), dtype=int)
    e = 0.5 * np.pi ** 2 * np.sum(n ** 2)
    for k in range(Mbi):
        e = MC_step(n, e, T)
    Emean[t] = e
    nsamples = 1
    for k in range(Mbi, M):
        Emean[t] += e
        nsamples += 1
    Emean[t] /= nsamples
plt.figure()
plt.plot(Tarr, Emean / N, '-o')
plt.plot(Tarr, 1.5 * Tarr, '--')
plt.xlabel('Temperature')
plt.ylabel('Mean energy')
plt.grid()
plt.show()