# Nicholas M Boffi
# 5/2/17
# 8.334 Final Project: Monte-Carlo Simulation of diffusion in an electron glass.
import numpy as np
np.set_printoptions(threshold='inf')
np.set_printoptions(linewidth='inf')
from scipy.constants import k as KB

def instantiate_grid(N, c):
    """
    Description:
    -----------
    Creates a random NxN grid of Ising spins correspond to electrons in holes.

    Input:
    ------
    N : The dimension of the square lattice.
    c : The concentration of particles (i.e., c*N*N = number of particles).

    Output:
    ------
    grid : The randomly initialized grid with the correct particle density.
    """

    # Create a grid of all holes.
    grid = -np.ones((N, N))

    # Determine the number of electrons.
    num_e = int(c*N*N)
    curr_num_e = 0

    # Fill the grid with this fraction.
    while (curr_num_e < num_e):
        x_ind = np.random.randint(N)
        y_ind = np.random.randint(N)

        if (grid[x_ind, y_ind] == -1):
            grid[x_ind, y_ind] = 1
            curr_num_e += 1

    return grid

def boltz_flip(grid, bonds, xx, yy, N, T):
    """
    Description:
    -----------
    Possibly the spin located at (xx, yy) according to standard Monte Carlo equilibration using the Boltzmann
    weight (rather than the Glauber dynamics) formalism to choose the probability.

    Input:
    ------
    grid  : The grid of spins.
    bonds : The array of bond energies.
    xx    : The x index of the particle to possibly flip.
    yy    : The y index of the particle to possibly flip.
    N     : The size of the N*N grid
    T     : The temperature at which the equilibration is performed.
    """ 
    # Determine if we are on the boundary, and if we need to handle the periodic
    # boundary conditions. Stands for right, left, top, and bottom boundary respectively.
    rb = (xx == N-1)
    lb = (xx == 0)
    tb = (yy == N-1)
    bb = (yy == 0)

    # Figure out the neighboring spin indices.
    # Right, left, up, down indices respectively.
    ri = 0   if rb else xx+1
    li = N-1 if lb else xx-1
    ui = 0   if tb else yy+1
    di = N-1 if bb else yy-1

    # Determine the negative of the energy change.
    dE = -2*grid[xx, yy]*(grid[li, yy]*bonds[li, yy] + grid[ri, yy]*bonds[ri, yy] + grid[xx, ui]*bonds[xx, ui] + grid[xx, di]*bonds[xx, di])

    # Determine if the flip happens.
    if (dE > 0): grid[xx, yy] = -grid[xx, yy]
    elif (np.random.uniform() < np.exp(dE/T)):
        grid[xx, yy] = -grid[xx, yy]

def step_forward(grid, bonds, N, T):
    """
    Description:
    -----------
    Picks a random spin, and attempts to flip it using the Boltzmann dynamics.

    Input:
    ------
    grid  : The grid of spin variables.
    bonds : The array of bond energies.
    N     : The size of the N*N square lattice.
    T     : The temperature at which the equilibration is performed.
    """
    ind = np.random.randint(0, N, 2)
    boltz_flip(grid, bonds, ind[0], ind[1], N, T)

def output_grid(grid):
    """
    Description:
    -----------
    Outputs the current state of the grid so the equilibration can be observed.

    Input:
    -----
    grid : The grid of spins.
    """
    # Simply print out the grid, using a 0 instead of -1 to represent vacancies for visibility.
    tmp = grid.copy()
    tmp[tmp == -1] = 0
    print tmp

def equilibrate(grid, bonds, N, T_start, T_final, num_steps, num_frames, output=False):
    """
    Description:
    ------------
    Equilibrates the grid via simulated annealing by performing num_steps flips starting at T_start and ending at T_final.
    The temperature is decreased by dt = (T_final - T_start)/num_steps each time.

    Input:
    ------
    grid       : The grid of spins.
    bonds      : The array of bond enregies.
    N          : The size of the N*N grid.
    T_start    : The initial equilibration temperature.
    T_final    : The final equilibration temperature.
    num_steps  : The number of spin flip attempts.
    num_frames : The number of times to output the grid.
    output     : Whether or not to print the grid.
    """ 
    # Determine when we output the grid and bond array.
    frame_step = num_steps / num_frames

    # Keep track of the current frame.
    frame = 0

    # Keep track of the simulated annealing current temperature.
    curr_T = T_start

    # Temperature increment so that we end on the final temperature at the last step.
    dT = (T_start - T_final)/num_steps

    # Do the update at each step.
    for ii in range(num_steps):
        # Print out the grid when necessary.
        if output and (ii % frame_step == 0):
            print 'Outputting grid on frame:', frame
            output_grid(grid)
            frame += 1

        # Keep lowering the temperature until we reach our final state.
        if (curr_T > T_final): curr_T -= dT

        # Flip a spin.
        step_forward(grid, bonds, N, curr_T)

def draw_bonds(N, Tm, ferro):
    """
    Description:
    ------------
    Draws and returns the array of random bond energies.
    Can draw from uniform or exponential distribution.

    Input:
    ------
    N     : Size of the N*N grid.
    Tm    : Maximum temperature the glass was exposed to (taken to be a parameter of the bond energy distribution).
    ferro : Boolean indicating whether or not the model is ferromagnetic.

    Output:
    -------
    mat : The array of random bond energies.
    """
    # Ensure that the matrix of bond energies is symmetric.
    mat = np.zeros((N, N))
    for ii in range(N):
        for jj in range(ii, N):
            # mat[ii, jj] = mat[jj, ii] = (np.random.exponential(Tm) + .05)
            mat[ii, jj] = mat[jj, ii] = Tm*np.random.uniform()

    # Ferromagnetic model.
    if ferro:
        return mat
    # Antiferromagnetic model.
    else:
        return -mat
 
def anneal(grid, bonds, N, T_start, T_final, num_anneals, num_frames, steps_per_spin, output=False):
    """
    Description:
    ------------
    'Anneals' the grid, by allowing it to equilibrate while cooling it from T_start to T_final num_anneals times.

    Input:
    ------
    grid           : The grid of spins.
    bonds          : The array of bond energies.
    T_start        : The initial temperature.
    T_final        : The final temperature.
    num_anneals    : Number of times to perform the equilibration.
    num_frames     : The number of times to output the grid per anneal.
    steps_per_spin : Number of MCS.
    output         : Whether or not to output the grid.

    """
    for run in range(num_anneals):
        print "Starting run:", run
        equilibrate(grid, bonds, N, T_start, T_final, steps_per_spin*N*N, num_frames, output)

if __name__ == '__main__':
    # Define the parameters.
    T_start = .75
    Tm = 1.
    T_final = Tm/5.
    N = 30
    num_anneals = 1
    num_frames = 1000
    steps_per_spin = 10000
    conc = .3
    ferro = True
    output = True

    # Create the grid and the bonds.
    grid  = instantiate_grid(N, conc)
    bonds = draw_bonds(N, Tm, ferro)

    # Anneal the grid.
    anneal(grid, bonds, N, T_start, T_final, num_anneals, num_frames, steps_per_spin, output)