# This is a simulation of the two-dimensional square-lattice Ising model
# written for Python 2.7. Input parameters are entered into the 
# metropolisAlgorithm call at the end of the file.
import random
import math


# couplingEnergy returns the coupling energy between two spins 
def couplingEnergy(spinA, spinB):
    return -spinA * spinB


# findLocalEnergy returns the total coupling energy between the specified spin
# site and its neighbors. If the isFlipped argument is True, the energy is
# calculated as though the spin had been flipped.
def findLocalEnergy(aLattice, aColumn, aRow, isFlipped):
    nColumns = len(aLattice)
    nRows = len(aLattice[0])
    if isFlipped:
        spinOfInterest = -aLattice[aColumn][aRow]
    else:
        spinOfInterest = aLattice[aColumn][aRow]
    neighborSpins = 0
    neighborSpins += aLattice[(aColumn - 1) % nColumns][aRow]
    neighborSpins += aLattice[(aColumn + 1) % nColumns][aRow]
    neighborSpins += aLattice[aColumn][(aRow - 1) % nRows]
    neighborSpins += aLattice[aColumn][(aRow + 1) % nRows]

    return -spinOfInterest * neighborSpins


# generateRandomLattice does exactly what it says on the tin. The function
# returns a list of lists of random 1 and -1 values. This is currently unused.
def generateRandomLattice(firstDimension, secondDimension):
    newLattice = []

    for i in range(firstDimension):
        newLattice.append([])
        for j in range(secondDimension):
            if(random.randint(0, 1) == 0):
                newLattice[i].append(-1)
            else:
                newLattice[i].append(1)
    return newLattice


# generateLattice returns a list of lists populated entirely by values of
# either 1 or -1.
def generateLattice(firstDimension, secondDimension):
    newLattice = []
    startingValue = random.choice([-1, 1])
    for i in range(firstDimension):
        newLattice.append([])
        for j in range(secondDimension):
            newLattice[i].append(startingValue)
    return newLattice


# checkSpinSite attempts to flip a spin according to the Metropolis algorithm.
# If a flip is performed, the function returns a list containing a new lattice,
# identical to the old one but with one spin flipped, the change in energy the
# flip yielded, and the new value of the spin.
def checkSpinSite(aLattice, aColumn, aRow, aTau):
    newEnergy = findLocalEnergy(aLattice, aColumn, aRow, True)
    oldEnergy = findLocalEnergy(aLattice, aColumn, aRow, False)
    energyDifference = newEnergy - oldEnergy

    if energyDifference < 0 or math.log(random.random()) < -energyDifference / aTau:
        aLattice[aColumn][aRow] = -aLattice[aColumn][aRow]
        return (aLattice, newEnergy, aLattice[aColumn][aRow])
    else:
        return (aLattice, oldEnergy, aLattice[aColumn][aRow])


# metropolisAlgorithm handles the majority of the program flow as well as the
# output. Its arguments are as follows:
# nColumns/nRows - the desired dimensions of the lattice.
# maxTemp - the maximum fundamental temperature of the lattice.
# tempStepSize - the amount the temperature is increased by each iteration.
# nEquilSweeps - the number of Metropolis sweeps the program executes to bring
# the lattice to thermal equilibrium before sampling can begin.
# nSamplingSweeps - the number of sampling sweeps at every temperature.
def metropolisAlgorithm(nColumns, nRows, maxTemp, tempStepSize, nEquilSweeps, nSampleSweeps):
    lattice = generateLattice(nColumns, nRows)
    tau = 0.001
    averageEnergies = []
    energyVariances = []
    averageMagnetizations = []
    magnetizationVariances = []
    taus = []
    while tau < maxTemp:
        print tau
        iSweep = 0
        while iSweep < nColumns * nRows * nEquilSweeps:
            randColumn = random.randint(0, nColumns - 1)
            randRow = random.randint(0, nRows - 1)

            step = checkSpinSite(lattice, randColumn, randRow, tau)
            if cmp(step[0], lattice):
                lattice = step[0]
            iSweep = iSweep + 1
        iSweep = 0
        averageEnergy = 0
        energyVariance = 0
        averageMagnetization = 0
        magnetizationVariance = 0
        energies = []
        magnetizations = []
        while iSweep < nColumns * nRows * nSampleSweeps:
            randColumn = random.randint(0, nColumns - 1)
            randRow = random.randint(0, nRows - 1)

            step = checkSpinSite(lattice, randColumn, randRow, tau)
            energies.append(step[1])
            magnetizations.append(step[2])
            averageEnergy = averageEnergy + step[1]
            averageMagnetization = averageMagnetization + step[2]
            if cmp(step[0], lattice):
                lattice = step[0]
            iSweep = iSweep + 1
        averageEnergy = averageEnergy / (nColumns * nRows * nSampleSweeps)
        averageEnergies.append(averageEnergy)
        averageMagnetization = averageMagnetization / (nColumns * nRows *
                                                       nSampleSweeps)
        averageMagnetizations.append(averageMagnetization / (nColumns * nRows *
                                                             nSampleSweeps))
        energyVariances.append(calculateVariance(energies))
        magnetizationVariances.append(calculateVariance(magnetizations))
        taus.append(tau)
        tau = tau + tempStepSize
    specificHeats = []
    maxSpecificHeat = 0
    specificHeatPeak = 0
    for i in range(len(energyVariances)):
        specificHeats.append(energyVariances[i] / taus[i] ** 2)
        if specificHeats[i] > maxSpecificHeat:
            maxSpecificHeat = specificHeats[i]
            specificHeatPeak = taus[i]
    print specificHeatPeak
    fileName = "simc" + str(nColumns) + "r" + str(nRows) + "t" + str(maxTemp) + "dt" + str(tempStepSize) + "eq" + str(nEquilSweeps) + "ss" + str(nSampleSweeps) + ".csv"
    outputToFile(open(fileName, "w"), [["Tau", taus], ["Ave Energy", averageEnergies], ["Energy Variance", energyVariances], ["Ave Magnetization", averageMagnetizations], ["Magnetization Variance", magnetizationVariances], ["Specific Heat", specificHeats]])


# Quick and dirty variance calculation. Inputting an empty list results in a
# return value of 0.
def calculateVariance(aList):
    summ = 0
    summSquared = 0
    n = 0
    for value in aList:
        n = n + 1
        summ = summ + value
        summSquared = summSquared + value * value
    if n > 0:
        return (summSquared - ((summ * summ) / n)) / n
    else:
        return 0


# outputToFile is just a small code snippet to cleanly prepare the csv to be
# written. argPairs must be of the format [["Value 1 Type", list1], ...]
def outputToFile(aFile, argPairs):
#    for pair in argPairs:
#        aFile.write(pair[0] + ",")
#    aFile.write("\n")
    for i in range(len(argPairs[0][1])):
        for j in range(len(argPairs)):
            aFile.write(str(argPairs[j][1][i]) + ",")
        aFile.write("\n")
    aFile.close()


metropolisAlgorithm(100, 100, 6.0, 0.01, 3.0, 1.0)