import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as opt
import scipy.special as spe

# Question 2:
def MC_sweep(theta, T, epsilon):
    N = theta.size
    for i in range(N):
        shift = np.random.random() # Generates a uniform random number between 0 and 1
        shift = - epsilon + 2. * epsilon * shift # Generates a uniform random number between -epsilon and epsilon
        sin_shift = np.sin(0.5 * shift) # sin(delta / 2)
        delta_E = (2. / N) * sin_shift * (np.sum(np.sin(theta[i] - theta + 0.5 * shift)) - sin_shift)
        if delta_E <= 0: # Metropolis
            theta[i] += shift
        elif T > 0. and np.random.random() < np.exp(- delta_E / T):
                theta[i] += shift
    
# Question 3: approximately 100 sweeps are required to reach the stationary state for all values of temperature. 
# This piece of code is to estimate the length of the burn-in phase.
N = 100
M = 5000
epsilon = 1.
T = .5
theta = np.random.uniform(- np.pi, np.pi, N)
S = np.zeros(M + 1)
S[0] = (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
for k in range(M):
   MC_sweep(theta, T, epsilon)
   S[k+1] = (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
plt.figure()
plt.plot(S)
plt.xlabel('Number of sweeps')
plt.ylabel(r'$S$')
plt.grid()
plt.show()

# For large T, S is close to 0, spins are not aligned.
# For smaller T, S increases, meaning that spins get aligned.
N = 100
Mbi = 100
M = 5000 # at least 50 times the length of the burn-in phase
epsilon = 1.
Tarr = np.arange(0., 1.1, 0.1)
Sarr = np.zeros_like(Tarr)
for t, T in enumerate(Tarr):
    theta = np.random.uniform(- np.pi, np.pi, N)
    for k in range(Mbi): 
       MC_sweep(theta, T, epsilon)
    Sarr[t] = (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
    ns = 1 # number of samples to compute the average
    for k in range(M):
       MC_sweep(theta, T, epsilon)
       Sarr[t] += (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
       ns += 1
    Sarr[t] /= ns
plt.figure()
plt.plot(Tarr, Sarr, '-o')
plt.xlabel('Temperature')
plt.ylabel(r'Order parameter')
plt.grid()
plt.show()

# Question 4: the curve becomes sharper when N icnreases, this signals the existence of a phase transition.
# In the thermodynamic limit (N->infinity), there is a phase transition for Tc=0.5.
Sarr = np.tile(Sarr, (3,1)) # We create an array of three rows for the three sizes 
Narr = np.array([20, 100, 200])
Mbi = 100
M = 5000 # at least 20 times the length of the burn-in phase
epsilon = 1.
plt.figure()
for n, N in enumerate(Narr):
    for t, T in enumerate(Tarr):
        print(N, T)
        theta = np.random.uniform(- np.pi, np.pi, N)
        for k in range(Mbi): 
           MC_sweep(theta, T, epsilon)
        Sarr[n,t] = (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
        ns = 1 # number of samples to compute the average
        for k in range(M):
           MC_sweep(theta, T, epsilon)
           Sarr[n,t] += (np.sum(np.cos(theta)) ** 2 + np.sum(np.sin(theta)) ** 2) / N ** 2
           ns += 1
        Sarr[n,t] /= ns
    plt.plot(Tarr, Sarr[n,:], '-o', label = str(N))
plt.xlabel('Temperature')
plt.ylabel(r'Order parameter')
plt.legend()
plt.grid()
plt.show()

# Question 5: for T=0.1, you can take epsilon between 1. and 2.7.
# The acceptance rate decreases when epsilon increases because we are trying to do bigger steps. 
def MC_sweep_with_acc_rate(theta, T, epsilon):
    global acc_rate
    N = theta.size
    for i in range(N):
        shift = np.random.random() # Generates a uniform random number between 0 and 1
        shift = - epsilon + 2. * epsilon * shift # Generates a uniform random number between -epsilon and epsilon
        sin_shift = np.sin(0.5 * shift) # sin(delta / 2)
        delta_E = (2. / N) * sin_shift * (np.sum(np.sin(theta[i] - theta + 0.5 * shift)) - sin_shift)
        if delta_E <= 0: # Metropolis
            theta[i] += shift
            acc_rate += 1.
        elif T > 0. and np.random.random() < np.exp(- delta_E / T):
                theta[i] += shift
                acc_rate += 1.
N = 100
M = 5000
T = 0.1
eps_arr = np.linspace(0., np.pi, 10)
acc_rate_arr = np.zeros_like(eps_arr)
for e, epsilon in enumerate(eps_arr):
    acc_rate = 0.
    theta = np.random.uniform(- np.pi, np.pi, N)
    for k in range(M): 
       MC_sweep_with_acc_rate(theta, T, epsilon)
    acc_rate_arr[e] = acc_rate / (M * N)
plt.figure()
plt.plot(eps_arr, acc_rate_arr * 100., '-o')
plt.xlabel(r'$\epsilon$')
plt.ylabel('Acceptance rate (%)')
plt.grid()
plt.show()