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

def calcVectorProjection(initialPos,pos,constraint):
	delta = pos - initialPos
	delta = delta - np.round(delta)
	return np.sum(delta*constraint)

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

stepSize = 2e-3 * np.sqrt(2048/nParticles)
maxF = 2 * np.sqrt(2048/nParticles)
phiString = '0.94'
xi = np.float64(xiString)
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) + '/' + xiString + '/' + 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(maxF/stepSize))
# p.setConstraintType(mc.enums.constraintEnum.stressControl)

if os.path.exists(saveDir + '/temp'):
	p.load(saveDir + '/initial')
	startingPos = p.getPositions()
	p.load(saveDir + '/temp')
	constraint = mc.fileIO.load2DArray(saveDir + '/constraint.dat',dtype=np.float64)
	stressStrain = mc.fileIO.load2DArray(saveDir + '/stressStrain.dat', dtype=np.float64)
	thisF = stressStrain[-1,1]
	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)
	startingPos = p.getPositions()
	p.save(saveDir + '/initial')
	stressStrain = np.zeros((numSteps,2))
	stressStrain[0,0] = calcVectorProjection(startingPos,p.getPositions(),constraint)
	stressStrain[0,1] = 0
	thisF = 0
	last=1

p.setConstraintType(mc.enums.constraintEnum.stressControl)
p.calcNeighborsCut(1)
p.calcForceEnergy()


if thisF < maxF:
	for i in range(last,numSteps):
		p.setExternalForce(stepSize*i*constraint)
		p.calcNeighborsCut(1)
		p.minimizeFIRE(criticalForce = critForce)
		stressStrain[i,0] = calcVectorProjection(startingPos,p.getPositions(),constraint)
		stressStrain[i,1] = stepSize * i
		thisF = stepSize * i
		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)