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])
xiString = sys.argv[4]
basedir = sys.argv[5]

# 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)

stepSize = 1e-4 / np.sqrt(2048/nParticles)
maxGamma = 0.1 / np.sqrt(2048/nParticles)
phiString = '0.94'
xi = np.float64(xiString)
critForce = 1e-12

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

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

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

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

saveFreq = 10
thisGamma = 0
numSteps = int(np.ceil(maxGamma/stepSize))
p.setConstraintType(mc.enums.constraintEnum.strainControl)

if os.path.exists(saveDir + '/temp'):
	p.load(saveDir + '/temp')
	constraint = mc.fileIO.load2DArray(saveDir + '/constraint.dat',dtype=np.float64)
	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()
	constraint = p.generateGCF(xi=xi)
	mc.fileIO.save2DArray(saveDir + '/constraint.dat',constraint)
	p.minimizeFIRE(criticalForce=critForce)
	p.save(saveDir + '/initial')
	stressStrain = np.zeros((numSteps,2))
	stressStrain[0,0] = 0
	stressStrain[0,1] = p.getConstraintMagnitude()*2.0*minRad
	thisGamma = 0
	last=1

p.setConstraintType(mc.enums.constraintEnum.strainControl)
p.setConstraint(constraint)
p.calcNeighborsCut(1)
p.calcForceEnergy()


if thisGamma < maxGamma:
	for i in range(last,numSteps):
		p.translate(stepSize,constraint)
		p.calcNeighborsCut(1)
		p.minimizeFIRE(criticalForce = critForce)
		stressStrain[i,0] = stepSize * i
		stressStrain[i,1] = p.getConstraintMagnitude()*2.0*minRad
		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)