import monteCarloPCP as mc
import numpy as np
import sys
import os
import time
import glob

dim = int(sys.argv[1])
nParticles = int(sys.argv[2])
potentialPower = int(sys.argv[3])
basedir = sys.argv[4]

stepSize = 1e-4
maxGamma = 0.2
phiString = '0.94'
critForce = 1e-12

# Set the random seed to generate GCF. To make reproducable, set seed to a saved integer
seed = np.random.randint(2**16)
p.setRandomSeed(seed)

p = mc.MCPacking(dim,nParticles,np.float64(phiString))

saveDir = basedir + '/' + str(dim) + '/' + str(nParticles) + '/AQS/' + str(seed) + '/'
os.makedirs(saveDir,exist_ok=True)

if dim is 2:
	p.setBiRadii()
else:
	p.setMonoRadii()

minRad = np.min(p.getRadii())
np.random.seed(seed)
p.setRandomSeed(seed)
p.setPotentialPower(potentialPower)

saveFreq = 10
thisGamma = 0
numSteps = int(np.ceil(maxGamma/stepSize))

if os.path.exists(saveDir + '/temp'):
	p.load(saveDir + '/temp')
	stressStrain = mc.fileIO.load2DArray(saveDir + '/stressStrain.dat', dtype=np.float64)
	thisGamma = stressStrain[-1,0]
	last = np.size(stressStrain,0)
	stressStrain = np.vstack((stressStrain,np.zeros((numSteps-last,2))))
else:
	p.setRandomPositions()
	p.calcNeighborsCut(1)
	p.calcForceEnergy()
	p.minimizeFIRE(criticalForce=critForce)
	p.save(saveDir + '/initial')
	stressStrain = np.zeros((numSteps,2))
	stressStrain[0,0] = 0
	stressStrain[0,1] = p.getStressTensor()[0,1]*(nParticles/p.getPhi())/((2.0*minRad)**2)
	thisGamma = 0
	last=1

p.setGeometryType(mc.enums.geometryEnum.leesEdwards)
p.calcNeighborsCut(1)
p.calcForceEnergy()


if thisGamma < maxGamma:
	for i in range(last,numSteps):
		p.setLEshift(stepSize*i)
		p.calcNeighborsCut(1)
		p.minimizeFIRE(criticalForce = critForce)
		stressStrain[i,0] = stepSize * i
		stressStrain[i,1] = p.getStressTensor()[0,1]*(nParticles/p.getPhi())/((2.0*minRad)**2)
		if (i%saveFreq==0 and i>0):
			p.save(saveDir + '/temp')
			mc.fileIO.save2DArray(saveDir + '/stressStrain.dat',stressStrain[:(i+1),:])

p.save(saveDir + '/temp')
mc.fileIO.save2DArray(saveDir + '/stressStrain.dat',stressStrain)