#!/usr/bin/python

# Parameters taken from: Comparison of phase-field and potts models for coarsening processes.
#  V. Tikare, E. A. Holm, D. Fan and L.Q. Chen, Acta Mat., 1998

import sys
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
from scipy.interpolate import griddata
import random as rd
import re

def PrintMsg(str,nx=None,ny=None):
    if nx is not None and ny is not None:
        print str,nx*ny
    else:
        print str

#Create_grid may be needed later so we are keeping it at the moment.
def create_grid(nx,ny):
    x=np.linspace(0,2*nx,nx+1)
    y=np.linspace(0,2*ny,ny+1)
    numCols = 2
    PrintMsg('Creating grid of size:',nx,ny)
    Grid=np.zeros(nx*ny*numCols).reshape(nx,ny,numCols)
    for i in range(0,nx):
        for j in range(0,ny):
            Grid[i,j,0]=x[i]
            Grid[i,j,1]=y[j]
    return Grid


# The main class definition for Phasefield.

def ReadRestart(filenam): #,numx,numy,RstState):
    f=open(filenam,'r')
    i=0
    Val=np.loadtxt(filenam)
    return Val.shape[0],Val.shape[1], Val

class McPotts():
    def __init__(self,nx,ny,numStates,restart=None):
        self.nx = nx
        self.ny = ny
        self.numStates=numStates
        self.Energy=np.zeros([self.nx,self.ny])
        if restart is None:
            self.State=np.random.randint(numStates,size=self.nx*self.ny).reshape(nx,ny)
        else:
            self.State = restart.astype(int)
#        self.__seeding()
        self.TotalEnergy = self.GetEnergy()



#Using random values for initialization leads to num overflows.
# From the Krill Paper, it was made clear that the initial values must be
# between -0.001 and 0.001 

#    def __seeding():
#        for i in range(0,self.nx):
#            for j in range(0,self.ny):
#                self.State=np.random.randint(0,self.numStates)
#


    def GetEnergy(self):
        ETotal = 0.0
        for i in range(0,self.nx):
            for j in range(0,self.ny):
                nNx1,nNx2,nPx1,nPx2 = self.NeighborIds(i)
                nNy1,nNy2,nPy1,nPy2 = self.NeighborIds(j)
                ETotal = ETotal + self.NeighborContribution(i,j,nNx1,j)+\
                         +  self.NeighborContribution(i,j,nPx1,j)+\
                         +  self.NeighborContribution(i,j,i,nNy1)+\
                         +  self.NeighborContribution(i,j,i,nPy1)+\
                         +  self.NeighborContribution(i,j,nNx2,j)+\
                         +  self.NeighborContribution(i,j,nPx2,j)+\
                         +  self.NeighborContribution(i,j,i,nNy2)+\
                         +  self.NeighborContribution(i,j,i,nPy2)

        return 0.5*ETotal


    def NeighborContribution(self,i,j,k,l):
        Val = 0.0
        if self.State[i,j] != self.State[k,l]:
            Val = 1.0
        return Val
                

# The system assumes periodic boundary conditions and therefore the contributions on
# grid points beyond the min and max index must be handled proprely.
    def NeighborIds(self,id):
        nN1 = id-1
        nN2 = id-2
        nP1 = id + 1
        nP2 = id + 2
        if id == 0:
            nN1 = self.nx-1
            nN2 = self.nx-2
        elif id == 1:
            nN2 = self.nx-1
        elif id == self.ny -1:
            nP1 = 0
            nP2 = 1
        elif id == self.ny-2:
            nP2 = 0
        return nN1, nN2, nP1,nP2

        
    def GrainGrowth(self,numSteps,Offset=0):
        for i in range(Offset+0,numSteps+Offset+3):
# Get random state
# Get random site, and assign previous state,
# Calculate total energy,
# Energy change > 0.0,  p = 0
# Energy change < 0.0  p = 1
# Generate random number between 0 and 1, if <= p, accept

            print 'Executing step number: ',i
            idx1= np.random.randint(0,self.nx)
            idx2= np.random.randint(0,self.ny)

            oldState = self.State[idx1,idx2]
            ind1= np.random.randint(0,self.numStates)
#            print ind1,oldState
            while ind1 == oldState:
                ind1 = np.random.randint(0,self.numStates)

            
            self.State[idx1,idx2] = ind1
            ENew=self.GetEnergy()

            prand = rd.random()
            ptest = 0
            if ENew -self.TotalEnergy <= 0:
                ptest = 1
            if prand <= ptest: # Im nachhinein thi is always true.
                print 'Accept: ',ENew
                self.TotalEnergy = ENew
            else:
                print 'Rejecting'
                self.State[idx1,idx2] = oldState
                
#            print self.State
#            plt.show()
            if i%1000 == 0:
                img=plt.imshow(self.State,cmap=cm.Greys_r)
                img.axes.get_xaxis().set_visible(False)
                img.axes.get_yaxis().set_visible(False)
                plt.savefig('./McPottsFigures/'+ 'McPottsGgrowth_'+str(i)+'.png',bbox_inches='tight')
                plt.close()
            if i%10000 == 0:
                filenam='restart_'+str(i)+'rst'
#                np.savetxt(filenam,(self.nx,self.ny))
#                np.savetxt(filenam,self.ny)
                strng=str(self.nx)+','+str(self.ny)
                np.savetxt(filenam,self.State,header=strng)

#            plt.show()

def main():
    if re.search('restart',sys.argv[1]):
        print 'Reading restart file name ',sys.argv[1]
        numx,numy,ReadRst = ReadRestart(sys.argv[1])
        Grains=McPotts(numx,numy,72,ReadRst)

    else:
        if len(sys.argv !=3):
            print 'Error '
            numx=int(sys.argv[1])
            numy=int(sys.argv[2])
        
        Grains = McPotts(numx,numy,72)
    
    Grains.GrainGrowth(5000000,1090000)


if __name__ == '__main__':
    main()