from yade import pack
from yade import export
from yade import timing
from yade import plot

if ('PFVFLOW' in features):

	num_spheres = 1000  # number of spheres
	young = 1e6
	compFricDegree = 3  # initial contact friction during the confining phase
	finalFricDegree = 30  # contact friction during the deviatoric loading
	mn, mx = Vector3(0, 0, 0), Vector3(1, 1, 0.4)  # corners of the initial packing
	graindensity = 2600
	toleranceWarning = 1.e-11
	toleranceCritical = 1.e-6

	O.materials.append(FrictMat(young=young, poisson=0.5, frictionAngle=radians(compFricDegree), density=graindensity, label='spheres'))
	O.materials.append(FrictMat(young=young, poisson=0.5, frictionAngle=0, density=0, label='walls'))
	walls = aabbWalls([mn, mx], thickness=0, material='walls')
	wallIds = O.bodies.append(walls)

	sp = pack.SpherePack()
	sp.makeCloud(mn, mx, -1, 0.3333, num_spheres, False, 0.95, seed=1)  #"seed" make the "random" generation always the same
	sp.toSimulation(material='spheres')

	triax = TriaxialStressController(
	        maxMultiplier=1. + 2e4 / young,  # spheres growing factor (fast growth)
	        finalMaxMultiplier=1. + 2e3 / young,  # spheres growing factor (slow growth)
	        thickness=0,
	        stressMask=7,
	        max_vel=0.005,
	        internalCompaction=True,  # If true the confining pressure is generated by growing particles
	)

	newton = NewtonIntegrator(damping=0.2)

	O.engines = [
	        ForceResetter(),
	        InsertionSortCollider([Bo1_Sphere_Aabb(), Bo1_Box_Aabb()]),
	        InteractionLoop(
	                [Ig2_Sphere_Sphere_ScGeom(), Ig2_Box_Sphere_ScGeom()], [Ip2_FrictMat_FrictMat_FrictPhys()], [Law2_ScGeom_FrictPhys_CundallStrack()],
	                label="iloop"
	        ),
	        FlowEngine(dead=1, label="flow"),  #introduced as a dead engine for the moment, see 2nd section
	        GlobalStiffnessTimeStepper(active=1, timeStepUpdateInterval=100, timestepSafetyCoefficient=0.8),
	        triax,
	        newton
	]

	triax.goal1 = triax.goal2 = triax.goal3 = -10000

	while 1:
		O.run(1000, True)
		unb = unbalancedForce()
		if unb < 0.001 and abs(-10000 - triax.meanStress) / 10000 < 0.001:
			break

	setContactFriction(radians(finalFricDegree))

	radius = 0
	for b in O.bodies:
		if b.state.mass == 0:
			b.state.blockedDOFs = 'xyzXYZ'
			b.state.vel = (0, 0, 0)
			b.state.angVel = (0, 0, 0)
		if b.state.mass > 0:
			radius += b.shape.radius
		#    b.state.blockedDOFs='xyz'
		#   b.state.vel=(0,0,0)
	radius = radius / num_spheres

	triax.dead = True
	while 1:
		O.run(1000, True)
		unb = unbalancedForce()
		if unb < 0.001:
			break

	press = 1000.
	O.run(10, 1)
	flow.dead = 0
	flow.meshUpdateInterval = -1
	flow.useSolver = 3
	flow.permeabilityFactor = 1
	flow.viscosity = 0.1
	flow.bndCondIsPressure = [0, 0, 1, 0, 0, 0]
	flow.bndCondValue = [0, 0, press, 0, 0, 0]
	flow.boundaryUseMaxMin = [0, 0, 0, 0, 0, 0]
	flow.iniVoidVolumes = True
	newton.damping = 0.1
	GlobalStiffnessTimeStepper.dead = True
	O.dt = min(0.8 * PWaveTimeStep(), 0.8 * 1. / 1200. * pi / flow.viscosity * graindensity * radius**2)
	O.dynDt = False

	O.run(1, 1)

	voidvol = 0.0
	voidvoltot = 0.0
	nvoids = flow.nCells()
	totalflux = [0] * (nvoids)
	initialvol = [0] * (nvoids)
	bar = [0] * (nvoids)
	initiallevel = O.bodies[flow.wallIds[flow.ymin]
	                       ].state.pos[1] + (O.bodies[flow.wallIds[flow.ymax]].state.pos[1] - O.bodies[flow.wallIds[flow.ymin]].state.pos[1]) / 3
	initialymax = O.bodies[flow.wallIds[flow.ymax]].state.pos[1]

	for ii in range(nvoids):
		initialvol[ii] = 1. / flow.getCellInvVoidVolume(ii)
		bar[ii] = flow.getCellBarycenter(ii)[1]
		voidvoltot += initialvol[ii]
		if bar[ii] >= initiallevel:
			voidvol += initialvol[ii]
			flow.setCellPImposed(ii, True)
			flow.setCellPressure(ii, 0)

	O.run(1, 1)
	iterini = O.time
	flow.saveVtk("./vtk", True)
	for ii in range(nvoids):
		if bar[ii] >= initiallevel:
			if flow.getCellPImposed(ii) == True:
				#totalflux[ii]=flow.getCellFluxFromId(ii)*O.dt
				totalflux[ii] = 0.0 * O.dt
	#Qx=(flow.getBoundaryFlux(flow.wallIds[flow.xmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.xmax]))*O.dt
	#Qz=(flow.getBoundaryFlux(flow.wallIds[flow.zmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.zmax]))*O.dt
	#Qin=(flow.getBoundaryFlux(flow.wallIds[flow.ymin]))*O.dt
	#Qout=(flow.getBoundaryFlux(flow.wallIds[flow.ymax]))*O.dt
	Qin = 0.0 * O.dt
	Qout = 0.0 * O.dt
	Qx = 0.0 * O.dt
	Qz = 0.0 * O.dt

	bubble = 0
	deltabubble = 0
	deltaovercross = 0
	deltaagain = 0

	def pressureImbibition():
		global bubble, deltabubble, deltaovercross, deltaagain
		global Qin, Qout, Qz, Qx, totalflux, total
		Qin += (flow.getBoundaryFlux(flow.wallIds[flow.ymin])) * O.dt
		Qout += (flow.getBoundaryFlux(flow.wallIds[flow.ymax])) * O.dt
		Qx += (flow.getBoundaryFlux(flow.wallIds[flow.xmin])) * O.dt + (flow.getBoundaryFlux(flow.wallIds[flow.xmax])) * O.dt
		Qz += (flow.getBoundaryFlux(flow.wallIds[flow.zmin])) * O.dt + (flow.getBoundaryFlux(flow.wallIds[flow.zmax])) * O.dt

		for ii in range(nvoids):
			if bar[ii] >= initiallevel:
				if flow.getCellPImposed(ii) == True:
					totalflux[ii] += flow.getCellFluxFromId(ii) * O.dt

		for ii in range(nvoids):
			if bar[ii] >= initiallevel:
				if flow.getCellPImposed(ii) == True:
					if (-totalflux[ii]) > initialvol[ii]:
						deltaflux = -totalflux[ii] - initialvol[ii]
						kk = 0
						neigh = [nvoids + 2, nvoids + 2, nvoids + 2, nvoids + 2]
						neighok = [nvoids + 2, nvoids + 2, nvoids + 2, nvoids + 2]
						for jj in range(4):
							if jj <= len(flow.getNeighbors(ii)) - 1:
								if flow.getCellPImposed(flow.getNeighbors(ii)[jj]) == True:
									if (-totalflux[flow.getNeighbors(ii)[jj]]) < initialvol[flow.getNeighbors(ii)[jj]]:
										neigh[kk] = flow.getNeighbors(ii)[jj]
										kk = kk + 1
						if kk == 0:
							totalflux[ii] += deltaflux
							bubble = bubble + 1
							deltabubble += deltaflux
						if kk != 0:
							totalflux[ii] += deltaflux
							deltafluxEach = deltaflux / kk
							deltadelta = 0
							for xx in range(4):
								if kk != 0:
									if neigh[xx] <= nvoids:
										jj = neigh[xx]
										if (-totalflux[jj] + deltafluxEach) > initialvol[jj]:
											deltadelta += -totalflux[jj] - initialvol[jj]
											totalflux[jj] += -(-totalflux[jj] - initialvol[jj])
											kk = kk - 1
											neighok[kk] = jj
							if kk == 0:
								deltaflux2 = deltaflux - deltadelta
								kk2 = 0
								neigh2 = [nvoids + 2] * (4 * 4)
								neighok2 = [nvoids + 2] * (4 * 4)
								for xx in range(4):
									if neigh[xx] <= nvoids:
										jj = neigh[xx]
										for gg in range(4):
											if gg <= len(flow.getNeighbors(jj)) - 1:
												if flow.getCellPImposed(flow.getNeighbors(jj)[gg]) == True:
													if (-totalflux[flow.getNeighbors(jj)[gg]]
													   ) < initialvol[flow.getNeighbors(jj)[gg]]:
														neigh2[kk2] = flow.getNeighbors(jj)[gg]
														kk2 = kk2 + 1
								if kk2 == 0:
									bubble = bubble + 1
									deltabubble += deltaflux2
								if kk2 != 0:
									deltafluxEach2 = deltaflux2 / kk2
									for xx in range(16):
										if kk2 != 0:
											if neigh2[xx] <= nvoids:
												gg = neigh2[xx]
												if (-totalflux[gg] + deltafluxEach2) > initialvol[gg]:
													print('check')
												if (-totalflux[gg] + deltafluxEach2) <= initialvol[gg]:
													deltadelta += deltafluxEach2
													totalflux[gg] += -deltafluxEach2
													kk2 = kk2 - 1
													neighok2[kk2] = gg
									if deltaflux != deltadelta:
										print('overcross')
										deltaovercross += +(deltaflux - deltadelta)
							if kk != 0:
								deltafluxEach2 = (deltaflux - deltadelta) / kk
								for xx in range(4):
									if kk != 0:
										if neigh[xx] <= nvoids:
											jj = neigh[xx]
											if jj != neighok[0] and jj != neighok[1] and jj != neighok[
											        2] and jj != neighok[3]:
												if (-totalflux[jj] + deltafluxEach2) <= initialvol[jj]:
													totalflux[jj] += -deltafluxEach2
													deltadelta += deltafluxEach2
													kk = kk - 1
													neighok[kk] = jj
								if deltaflux != deltadelta:
									print('again')
									deltaagain += +(deltaflux - deltadelta)

		total = 0
		for ii in range(nvoids):
			total += totalflux[ii]
			if bar[ii] >= initiallevel:
				if flow.getCellPImposed(ii) == True:
					if -(totalflux[ii]) == initialvol[ii]:
						flow.setCellPImposed(ii, False)
						#print 'pp_',ii
					if -(totalflux[ii]) > initialvol[ii]:
						flow.setCellPImposed(ii, False)
						print('error_', ii)
		total = abs(total)

	#  for ii in range(nvoids):
	#    if bar[ii]>=initiallevel:

	#        if flow.getCellPImposed(ii)==False:
	#          if -(totalflux[ii])!=initialvol[ii]:
	#            print 'error_',ii

	file = open('Test.txt', "w")
	checkdifference = 0

	def equilibriumtest():
		global F33, F22, checkdifference, errors
		#unbalanced=unbalancedForce()
		F33 = abs(O.forces.f(flow.wallIds[flow.ymax])[1])
		F22 = abs(O.forces.f(flow.wallIds[flow.ymin])[1])
		#F11 =abs(O.forces.f(flow.wallIds[flow.xmax])[0]),
		#F00=abs(O.forces.f(flow.wallIds[flow.xmin])[0]),
		#F44=abs(O.forces.f(flow.wallIds[flow.zmin])[2]),
		#F55=abs(O.forces.f(flow.wallIds[flow.zmax])[2]),
		deltaF = abs(F33 - F22)
		file.write(str(O.iter) + " " + str(deltaF) + "\n")
		if O.time >= iterini + 1.5:
			if checkdifference == 0:
				print('check F done')
				if deltaF > 0.01 * press:
					raise YadeCheckError('Error: too high difference between forces acting at the bottom and upper walls')
					#O.pause()
				checkdifference = 1

	once = 0

	def fluxtest():
		global once, QinOk
		no = 0
		QinOk = -Qin - deltabubble
		error = QinOk - total
		if error > toleranceWarning:
			print(
			        "Warning: difference between total water volume flowing through bottom wall and water loss due to air bubble generations",
			        QinOk, " vs. total water volume flowing inside dry or partially saturated cells", total
			)
		if error > toleranceCritical:
			raise YadeCheckError("The difference is more, than the critical tolerance!")
		file.write(str(O.time) + " " + str(QinOk) + " " + str(error) + " \n")
		for ii in range(nvoids):
			if bar[ii] >= initiallevel:
				if flow.getCellPImposed(ii) == True:
					no = 1
		if once == 0:
			if no == 0:
				imbtime = O.time - iterini
				print(imbtime, voidvol, total, QinOk)
				if voidvol - total > toleranceWarning:
					print(
					        "Warning: initial volume of dry voids", voidvol,
					        " vs. total water volume flowing inside dry or partially saturated cells", total
					)
				if voidvol - total > toleranceCritical:
					raise YadeCheckError("The difference is more, than the critical tolerance!")
				file.write(str(imbtime) + " " + str(voidvol) + " " + str(total) + " " + str(QinOk) + "\n")
				once = 1
				timing.stats()

	def addPlotData():
		global F33, F22, QinOk, total
		plot.addData(i1=O.iter, t=O.time, Fupper=F33, Fbottom=F22, Q=QinOk, T=total)

	plot.live = True
	plot.plots = {'t': ('Fbottom', 'Fupper'), ' t ': ('Q', 'T')}
	plot.plot()

	def pl():
		flow.saveVtk("./vtk", True)

	O.engines = O.engines + [PyRunner(iterPeriod=100, command='pl()')]
	#O.engines=O.engines+[VTKRecorder(iterPeriod=100,recorders=['spheres'],fileName='./exp')]
	O.engines = O.engines + [PyRunner(iterPeriod=1, command='equilibriumtest()')]
	O.engines = O.engines + [PyRunner(iterPeriod=1, command='pressureImbibition()')]
	O.engines = O.engines + [PyRunner(iterPeriod=1, command='fluxtest()')]
	O.engines = O.engines + [PyRunner(iterPeriod=1, command='addPlotData()')]

	O.timingEnabled = True
	#file.close()
	#plot.saveDataTxt('plots.txt',vars=('i1','t','Fupper','Fbottom','Q','T'))

	import tempfile, shutil
	dirpath = tempfile.mkdtemp()
	for fileName in ['./vtk', './Test.txt']:
		if (os.path.exists(fileName)):
			shutil.move(fileName, dirpath)
		print("File %s moved into %s/ directory" % (fileName, dirpath))

else:
	print("This checkDEM-PFVPressureInjection.py cannot be executed because PFVFLOW is disabled")
