# -*- encoding=utf-8 -*-
################################################################################
# Load the saved confined cubic aggregate using a KUBC (Kinematic Uniform
# Boundary Condition) applied as a velocity gradient tensor on the aggregate
# power diagram described by the alpha shape surface reconstruction algorithm
################################################################################
import os, sys, gts
from scipy.interpolate import griddata
from yade import pack, export
import numpy as np


#######################################
###       FUNCTION DEFINITIONS      ###
#######################################
def averageDefGrad(xq, Aq, vol0):  # Can be used for Deformation/Velocity Gradient Tensors
	defGrad = Matrix3.Zero
	for i, el in enumerate(Aq):
		defGrad = defGrad + Vector3.outer(xq[i], Aq[i]) / vol0
	return defGrad


def triVelMicro(strainPt, triThresh):
	global rate_v, currF
	currStrain = currF - Matrix3.Identity
	deltaStrain = strainPt - currStrain  #strainPath[ptNum]
	strainPtSum, startSum, deltaSum = strainPt.sum(), currStrain.sum(), deltaStrain.sum()
	print(startSum, deltaSum, strainPtSum)
	rate_v = (deltaStrain * currF.inverse()).normalized() * boundVel
	stage = 1  # Loading stages 1) Increment Strain 2) Relaxation
	appCoeff = 20.  # Approach Coefficient for decreasing Velocity Gradient near target strain
	meanStress = TW.calcAlphaStress(alpha, alphaShrinked, True)
	while 1:
		oldStress = meanStress
		for j in range(40):
			TW.applyAlphaVel(rate_v, alpha, alphaShrinked, True)
			t0 = O.time
			O.run(5, True)
			if stage == 1:
				Fdot = rate_v * currF
				currF += (Fdot) * (O.time - t0)
				currStrain = currF - Matrix3.Identity
				currSum = currStrain.sum()
			unb = unbalancedForce()
			meanStress = TW.calcAlphaStress(alpha, alphaShrinked, True)
			triRatio = 1. - abs(currSum - startSum) / abs(deltaSum)
			vMag = boundVel * max(0.005, min(1.0, triRatio * appCoeff))
			rate_tmp = deltaStrain * currF.inverse()
			rate_v = rate_tmp.normalized() * vMag if stage == 1 else Matrix3.Zero
		sflag = 0
		for i in range(0, 3):
			for j in range(0, 3):
				triDiff = abs((meanStress[i][j] - oldStress[i][j]) / max(abs(meanStress[i][j]), 0.2 * meanStress.norm()))
				if triDiff > triThresh:
					sflag += 1
		print('unb:', unb, ' stress flags', sflag)
		print('Percent of Strain Target', 100. * abs(currSum - startSum) / abs(deltaSum))
		print('Stress Tensor\n', np.asarray(meanStress))
		sys.stdout.flush()
		if sflag == 0 and stage == 2:
			break
		if abs(currSum - startSum) >= abs(deltaSum) and stage == 1:
			stage = 2


def cart2sph(cartX, cartY, cartZ):
	point = Vector3(cartX, cartY, cartZ)
	sphR = np.sqrt(point.dot(point))
	sphTheta = np.arctan2(point[1], point[0])
	sphPhi = np.arccos(point[2] / sphR)
	return sphTheta, sphPhi, sphR


def sph2cart(theta, phi, r):
	x = r * np.cos(theta) * np.sin(phi)
	y = r * np.sin(theta) * np.sin(phi)
	z = r * np.cos(phi)
	return x, y, z


def ancillaryGrid(alpha, alphaShrinked, isFix, s, newBound):
	TW = TesselationWrapper()
	TW.triangulate()
	segments = TW.getAlphaGraph(alpha, alphaShrinked, isFix)
	caps = TW.getAlphaCaps(alpha, alphaShrinked, isFix)
	# aaBB Center and Size
	center, length = np.mean(aabbExtrema(), axis=0), np.mean(aabbDim(), axis=0)
	# Use alphaGraph and alphaCaps to build a boundary grid
	polygons = np.asarray([cart2sph(*vertex - center) for vertex in segments])
	polygons = np.vstack([tuple(row) for row in polygons])
	boundGrid1 = np.vstack([polygons])
	# Mirror points to create a continuous surface in spherical coords
	boundGrid = np.vstack([boundGrid1 + [revTh, revPh, 0.] for revTh in [-2 * np.pi, 0, 2 * np.pi] for revPh in [-np.pi, 0, np.pi]])
	# Convert unit cube to spherical coordinates and interpolate surface
	gridGts = np.asarray([cart2sph(*vert) for vert in newBound])
	gridGts[:, 2] = griddata(boundGrid[:, 0:2], boundGrid[:, 2], gridGts[:, 0:2])
	for index, vertex in enumerate(s.vertices()):
		vertex.x, vertex.y, vertex.z = (sph2cart(*gridGts[index]) + center)
	return [
	        [Vector3(np.mean([vertex.coords() for vertex in face.vertices()], axis=0)) for face in s.faces()],
	        [Vector3(face.normal()) / 2. for face in s.faces()],
	        s.volume()
	]


def sepath(alpha, alphaShrinked, isFix):
	global currF
	print('********** Saving State Variables ***************')
	newBound = [ogPt * currF for ogPt in ogBound]
	bSurf.append(ancillaryGrid(alpha, alphaShrinked, isFix, unitCube, newBound))
	print('currF\n', np.asarray(currF))
	Fbar.append(averageDefGrad(bSurf[-1][0], bSurf[0][1], bSurf[0][2]))
	print('Fbar\n', np.asarray(Fbar[-1]))
	# Note that currF and Fbar are both evaluations of the deformation gradient
	# currF tracks from velocity gradient, while Fbar obtains it after reinterpolating
	# the surface grid to the latest power diagram
	Sbar.append(TW.calcAlphaStress(alpha, alphaShrinked, True))
	FbarRef, FbarCurr, SbarRef, SbarCurr = Fbar[0], Fbar[-1], Sbar[0], Sbar[-1]
	eGreen = 0.5 * (FbarCurr * FbarCurr.transpose() - Matrix3.Identity)
	hydroStress = SbarCurr.trace() / 3.e6
	dhydroStress = (SbarCurr.trace() - SbarRef.trace()) / 3.e6
	volStrain = FbarCurr.determinant() - 1.
	pStress = SbarCurr.spectralDecomposition()[1] / 1.e6 - SbarRef.spectralDecomposition()[1] / 1.e6
	deviatorStress = pStress.maxCoeff() - pStress.minCoeff()
	axStrain = eGreen[0][0]
	pStrain = eGreen.spectralDecomposition()[1]
	deviatorStrain = pStrain.maxCoeff() - pStrain.minCoeff()
	with open('resultsKUBC/cauchygreen.txt', 'a') as f:
		f.write('%f %f %f %f %f\n' % (hydroStress, deviatorStress, volStrain, axStrain, deviatorStrain))


def vtkOut(suffix, alpha, alphaShrinked, isFix):
	vtkExporter = yade.export.VTKExporter('resultsKUBC/spheres' + suffix)
	vtkExporter.exportSpheres()

	TW.triangulate()
	segments = TW.getAlphaGraph(alpha, alphaShrinked, True)
	caps = TW.getAlphaCaps(alpha, alphaShrinked, True)

	f = open('resultsKUBC/LVdiagram' + suffix + '.vtk', 'w')
	f.write("# vtk DataFile Version 2.0\nGenerated from YADE\n")
	f.write("ASCII\nDATASET POLYDATA\n")
	f.write("POINTS %d float \n" % (len(segments)))
	[f.write("%f %f %f\n" % (point[0], point[1], point[2])) for point in segments]
	f.write("LINES %d %d \n" % (len(segments) / 2, 3 * len(segments) / 2))
	[f.write("2 %d %d\n" % (2 * k, 2 * k + 1)) for k in range(int(len(segments) / 2))]
	f.close()

	f = open('resultsKUBC/gtsSurface' + suffix + '.vtk', 'w')
	unitCube.write_vtk(f)
	f.close()


############################################
###     LOAD REFERENCE CONFIGURATION     ###
############################################
if not (os.path.isdir('resultsKUBC')):
	os.mkdir('resultsKUBC')

O.load('confinedState.yade.gz')
O.resetTime()
############################################
###           FIND ALPHA GRAPH           ###
############################################
minRad = min(b.shape.radius for b in O.bodies)
meanDia = 2. * sum([b.shape.radius for b in O.bodies]) / len(O.bodies)
alpha = (22. * minRad)**2
alphaShrinked = (21. * minRad)**2
isFix = True

TW = TesselationWrapper()
TW.triangulate()
############################################
###         INITIALIZE VARIABLES         ###
############################################
stressM, rate_v = Matrix3.Identity * (-5.0e6), Matrix3.Zero
TW.applyAlphaForces(stressM, alpha, alphaShrinked, True)
O.step()
Sbar = [TW.calcAlphaStress(alpha, alphaShrinked, True)]
print("Initial Stress Tensor \n", np.asarray(Sbar[0]))
currF = Matrix3.Identity
Fbar = [currF]
with open('resultsKUBC/cauchygreen.txt', 'w') as f:
	f.write('pStress qStress vStrain aStrain dStrain\n')
###########################################
###      ANCILLARY GRID AT REF. STATE   ###
###########################################
# Build a unit gts cube and subdivide it 'subs' times
subs = 5
unitCube = gts.cube()
for i in range(subs):
	unitCube.tessellate()
	ogBound = [Vector3(vertex.coords()) for vertex in unitCube.vertices()]
	for i, vertex in enumerate(unitCube.vertices()):
		vertex.x, vertex.y, vertex.z = (ogBound[i] / ogBound[i].maxAbsCoeff())
# Use the unit gts cube to interpolate ancillary surface grid
ogBound = [Vector3(vertex.coords()) for vertex in unitCube.vertices()]
bSurf = []
sepath(alpha, alphaShrinked, True)
###########################################
###      TRIAXIAL STRESS/STRAIN GOALS   ###
###########################################
# Inertial Number
iNum = 7.9e-5
# Rate of Deformation Tensor Increment
boundVel = .45 * iNum * sqrt(5.e6 / O.bodies[0].mat.density) / meanDia
print('Quasistatic Strain Rate', boundVel)
###########################################
###             LOADING PATH            ###
###########################################
strainInc = -0.001
strainPath = [Matrix3.Zero]
# Load by an isotropic strain increment
strainPath.append(strainPath[-1] + Matrix3.Identity * strainInc)
# # Load by a uniaxial strain increment
# strainPath.append(strainPath[-1] + Matrix3(Vector3(strainInc,0.,0.),Vector3.Zero, Vector3.Zero))
# # Load by a deviatoric strain increment
# strainPath.append(strainPath[-1] + Matrix3(Vector3(0.,0.,0.5*strainInc),
#                                     Vector3.Zero,
#                                     Vector3(0.5*strainInc,0.,0.)))
###########################################
###            BOUNDARY CONDITION       ###
###########################################
print('APPLY KUBC')
for ptNum, strainPt in enumerate(strainPath[1:]):
	print('Apply a strain of \n', np.asarray(strainPt))
	stepper.timestepSafetyCoefficient = 0.6
	unbTol, stressTol = 0.01, 0.005
	loader.iterPeriod = 0  # Loading is done in function triVelMicro()
	triVelMicro(strainPt, stressTol)
	sepath(alpha, alphaShrinked, True)
######################################
###             VTK EXPORTS        ###
######################################
vtkOut("KUBC", alpha, alphaShrinked, isFix)
