# -*- encoding=utf-8 -*-
from __future__ import print_function
################################################################################
#
# Triaxial test. Axial strain rate is prescribed and transverse prestress.
# Test is possible on prism or cylinder
# An independent c++ engine may be created from this script in the future.
#
################################################################################
from builtins import range
from yade import pack, plot
import os

# default parameters or from table
readParamsFromTable(
        noTableOk=True,
        # type of test ['cyl','cube']
        testType='cyl',

        # material parameters
        young=20e9,
        poisson=.2,
        frictionAngle=1.2,
        sigmaT=1.5e6,
        epsCrackOnset=1e-4,
        relDuctility=30,

        # prestress
        preStress=-3e6,
        # axial strain rate
        strainRate=-100,

        # assamlby parameters
        rParticle=.075e-3,  #
        width=2e-3,
        height=5e-3,
        bcCoeff=5,

        # facets division
        nw=24,
        nh=15,

        # output specifications
        fileName='test',
        exportDir='/tmp',
        runGnuplot=False,
        runInGui=True,
)
from yade.params.table import *
assert testType in ['cyl', 'cube']

# materials
concMat = O.materials.append(
        CpmMat(young=young, frictionAngle=frictionAngle, poisson=poisson, sigmaT=sigmaT, epsCrackOnset=epsCrackOnset, relDuctility=relDuctility)
)
frictMat = O.materials.append(FrictMat(young=young, poisson=poisson, frictionAngle=frictionAngle))

# spheres
pred = pack.inCylinder((0, 0, 0), (0, 0, height), .5 * width) if testType == 'cyl' else pack.inAlignedBox(
        (-.5 * width, -.5 * width, 0), (.5 * width, .5 * width, height)
) if testType == 'cube' else None
sp = SpherePack()
sp = pack.randomDensePack(pred, spheresInCell=2000, radius=rParticle, memoizeDb='/tmp/triaxTestOnCylinder.sqlite', returnSpherePack=True)
spheres = sp.toSimulation(color=(0, 1, 1), material=concMat)

# bottom and top of specimen. Will have prescribed velocity
bot = [O.bodies[s] for s in spheres if O.bodies[s].state.pos[2] < rParticle * bcCoeff]
top = [O.bodies[s] for s in spheres if O.bodies[s].state.pos[2] > height - rParticle * bcCoeff]
vel = strainRate * (height - rParticle * 2 * bcCoeff)
for s in bot:
	s.shape.color = (1, 0, 0)
	s.state.blockedDOFs = 'xyzXYZ'
	s.state.vel = (0, 0, -vel)
for s in top:
	s.shape.color = Vector3(0, 1, 0)
	s.state.blockedDOFs = 'xyzXYZ'
	s.state.vel = (0, 0, vel)

# facets
facets = []
if testType == 'cyl':
	rCyl2 = .5 * width / cos(pi / float(nw))
	for r in range(nw):
		for h in range(nh):
			v1 = Vector3(rCyl2 * cos(2 * pi * (r + 0) / float(nw)), rCyl2 * sin(2 * pi * (r + 0) / float(nw)), height * (h + 0) / float(nh))
			v2 = Vector3(rCyl2 * cos(2 * pi * (r + 1) / float(nw)), rCyl2 * sin(2 * pi * (r + 1) / float(nw)), height * (h + 0) / float(nh))
			v3 = Vector3(rCyl2 * cos(2 * pi * (r + 1) / float(nw)), rCyl2 * sin(2 * pi * (r + 1) / float(nw)), height * (h + 1) / float(nh))
			v4 = Vector3(rCyl2 * cos(2 * pi * (r + 0) / float(nw)), rCyl2 * sin(2 * pi * (r + 0) / float(nw)), height * (h + 1) / float(nh))
			f1 = facet((v1, v2, v3), color=(0, 0, 1), material=frictMat)
			f2 = facet((v1, v3, v4), color=(0, 0, 1), material=frictMat)
			facets.extend((f1, f2))
elif testType == 'cube':
	nw2 = nw / 4
	for r in range(nw2):
		for h in range(nh):
			v11 = Vector3(-.5 * width + (r + 0) * width / nw2, -.5 * width, height * (h + 0) / float(nh))
			v12 = Vector3(-.5 * width + (r + 1) * width / nw2, -.5 * width, height * (h + 0) / float(nh))
			v13 = Vector3(-.5 * width + (r + 1) * width / nw2, -.5 * width, height * (h + 1) / float(nh))
			v14 = Vector3(-.5 * width + (r + 0) * width / nw2, -.5 * width, height * (h + 1) / float(nh))
			f11 = facet((v11, v12, v13), color=(0, 0, 1), material=frictMat)
			f12 = facet((v11, v13, v14), color=(0, 0, 1), material=frictMat)
			v21 = Vector3(+.5 * width, -.5 * width + (r + 0) * width / nw2, height * (h + 0) / float(nh))
			v22 = Vector3(+.5 * width, -.5 * width + (r + 1) * width / nw2, height * (h + 0) / float(nh))
			v23 = Vector3(+.5 * width, -.5 * width + (r + 1) * width / nw2, height * (h + 1) / float(nh))
			v24 = Vector3(+.5 * width, -.5 * width + (r + 0) * width / nw2, height * (h + 1) / float(nh))
			f21 = facet((v21, v22, v23), color=(0, 0, 1), material=frictMat)
			f22 = facet((v21, v23, v24), color=(0, 0, 1), material=frictMat)
			v31 = Vector3(+.5 * width - (r + 0) * width / nw2, +.5 * width, height * (h + 0) / float(nh))
			v32 = Vector3(+.5 * width - (r + 1) * width / nw2, +.5 * width, height * (h + 0) / float(nh))
			v33 = Vector3(+.5 * width - (r + 1) * width / nw2, +.5 * width, height * (h + 1) / float(nh))
			v34 = Vector3(+.5 * width - (r + 0) * width / nw2, +.5 * width, height * (h + 1) / float(nh))
			f31 = facet((v31, v32, v33), color=(0, 0, 1), material=frictMat)
			f32 = facet((v31, v33, v34), color=(0, 0, 1), material=frictMat)
			v41 = Vector3(-.5 * width, +.5 * width - (r + 0) * width / nw2, height * (h + 0) / float(nh))
			v42 = Vector3(-.5 * width, +.5 * width - (r + 1) * width / nw2, height * (h + 0) / float(nh))
			v43 = Vector3(-.5 * width, +.5 * width - (r + 1) * width / nw2, height * (h + 1) / float(nh))
			v44 = Vector3(-.5 * width, +.5 * width - (r + 0) * width / nw2, height * (h + 1) / float(nh))
			f41 = facet((v41, v42, v43), color=(0, 0, 1), material=frictMat)
			f42 = facet((v41, v43, v44), color=(0, 0, 1), material=frictMat)
			facets.extend((f11, f12, f21, f22, f31, f32, f41, f42))
O.bodies.append(facets)
mass = O.bodies[0].state.mass
for f in facets:
	f.state.mass = mass
	f.state.blockedDOFs = 'XYZz'

# plots
plot.plots = {
        'e': ('s',),
}


def plotAddData():
	f1 = sum(O.forces.f(b.id)[2] for b in top)
	f2 = sum(O.forces.f(b.id)[2] for b in bot)
	f = .5 * (f2 - f1)
	s = f / (pi * .25 * width * width) if testType == 'cyl' else f / (width * width) if testType == 'cube' else None
	e = (top[0].state.displ()[2] - bot[0].state.displ()[2]) / (height - rParticle * 2 * bcCoeff)
	plot.addData(
	        i=O.iter,
	        s=s,
	        e=e,
	)


# apply prestress to facets
def addForces():
	for f in facets:
		n = f.shape.normal
		a = f.shape.area
		O.forces.addF(f.id, preStress * a * n)


# stop condition and exit of the simulation
def stopIfDamaged(maxEps=5e-3):
	extremum = max(abs(s) for s in plot.data['s'])
	s = abs(plot.data['s'][-1])
	e = abs(plot.data['e'][-1])
	if O.iter < 1000 or s > .5 * extremum and e < maxEps:
		return
	f = os.path.join(exportDir, fileName)
	print('gnuplot', plot.saveGnuplot(f, term='png'))
	if runGnuplot:
		import subprocess
		os.chdir(exportDir)
		subprocess.Popen(['gnuplot', f + '.gnuplot']).wait()
	print('Simulation finished')
	O.pause()
	#sys.exit(0) # results in some threading exception


O.dt = .5 * utils.PWaveTimeStep()
enlargeFactor = 1.5
O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=enlargeFactor, label='bo1s'),
                               Bo1_Facet_Aabb()]),
        InteractionLoop(
                [
                        Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=enlargeFactor, label='ss2d3dg'),
                        Ig2_Facet_Sphere_ScGeom(),
                ],
                [
                        Ip2_CpmMat_CpmMat_CpmPhys(cohesiveThresholdIter=O.iter + 5),
                        Ip2_FrictMat_CpmMat_FrictPhys(),
                        Ip2_FrictMat_FrictMat_FrictPhys(),
                ],
                [
                        Law2_ScGeom_CpmPhys_Cpm(),
                        Law2_ScGeom_FrictPhys_CundallStrack(),
                ],
        ),
        PyRunner(iterPeriod=1, command="addForces()"),
        NewtonIntegrator(damping=.3),
        CpmStateUpdater(iterPeriod=50, label='cpmStateUpdater'),
        PyRunner(command='plotAddData()', iterPeriod=10),
        PyRunner(iterPeriod=50, command='stopIfDamaged()'),
]

# run one step
O.step()

# reset interaction detection enlargement
bo1s.aabbEnlargeFactor = ss2d3dg.interactionDetectionFactor = 1.0

# initialize auto-updated plot
if runInGui:
	plot.plot()
	try:
		from yade import qt
		renderer = qt.Renderer()
		# uncomment following line to exagerate displacement
		#renderer.dispScale=(100,100,100)
	except:
		pass

# run
O.run()
