# -*- encoding=utf-8 -*-
#*************************************************************************
#  Copyright (C) 2019 by Robert Caulk                                    *
#  rob.caulk@gmail.com                                                   *
#                                                                        *
#  This program is free software; it is licensed under the terms of the  *
#  GNU General Public License v2 or later. See file LICENSE for details. *
#*************************************************************************/
#
# Script demonstrating the use of ThermalEngine by comparing conduction
# scheme to analytical solution to Fourier (rod cooling with constant 
# boundary conditions). See details in:
#
# Caulk, R., Scholtes, L., Kraczek, M., Chareyre, B. (In Print) A 
# pore-scale Thermo-Hydro-Mechanical coupled model for particulate systems. 
# Computer Methods in Applied Mechanics and Engineering. Accepted July 2020.
#

from yade import pack
from yade import timing
import numpy as np
import shutil
timeStr = time.strftime('%m-%d-%Y')
num_spheres=1000# number of spheres
young=1e6
rad=0.003

mn,mx=Vector3(0,0,0),Vector3(1.0,0.008,0.008) # corners of the initial packing

thermalCond = 2. #W/(mK)
heatCap = 710. #J(kg K) 
t0 = 400. #K

r = rad
k = 2*2.0*r  # 2*k*r 
Cp = 710.
rho = 2600.
D = 2.*r
m = 4./3.*np.pi*r**2/rho
# macro diffusivity
thermalDiff = 6.*k/(D*np.pi*Cp*rho)

identifier = '-conductionVerification'

if not os.path.exists('VTK'+timeStr+identifier):
	os.mkdir('VTK'+timeStr+identifier)
else:
	shutil.rmtree('VTK'+timeStr+identifier)
	os.mkdir('VTK'+timeStr+identifier)

if not os.path.exists('txt'+timeStr+identifier):
	os.mkdir('txt'+timeStr+identifier)
else:
	shutil.rmtree('txt'+timeStr+identifier)
	os.mkdir('txt'+timeStr+identifier)

shutil.copyfile(sys.argv[0],'txt'+timeStr+identifier+'/'+sys.argv[0])

O.materials.append(FrictMat(young=young,poisson=0.5,frictionAngle=radians(3),density=2600,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)

O.bodies.append(pack.regularOrtho(pack.inAlignedBox(mn,mx),radius=rad,gap=-1e-8,material='spheres'))

print('num bodies ', len(O.bodies))

ThermalEngine = ThermalEngine(dead=1,label='thermal');

newton=NewtonIntegrator(damping=0.2)
intRadius = 1
O.engines=[
	ForceResetter(),
	InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=intRadius),Bo1_Box_Aabb()]),
	InteractionLoop(
		[Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=intRadius),Ig2_Box_Sphere_ScGeom()],
		[Ip2_FrictMat_FrictMat_FrictPhys()],
		[Law2_ScGeom_FrictPhys_CundallStrack()],label="iloop"
	),
	FlowEngine(dead=1,label="flow",multithread=False),
	ThermalEngine,	GlobalStiffnessTimeStepper(active=1,timeStepUpdateInterval=100,timestepSafetyCoefficient=0.8),
	#triax,
	VTKRecorder(iterPeriod=500,fileName='VTK'+timeStr+identifier+'/spheres-',recorders=['spheres','thermal','intr'],dead=1,label='VTKrec'),
	newton
]

#goal = -1e5
#triax.goal1=triax.goal2=triax.goal3=goal

for b in O.bodies:
	if isinstance(b.shape, Sphere): 
		b.dynamic=False


# we only need flow engine to detect boundaries, there is no flow computed for this 
flow.dead=0

thermal.dead=0
thermal.conduction=True
thermal.thermoMech=False
thermal.advection=False
thermal.fluidThermoMech = False
thermal.solidThermoMech = False
thermal.fluidConduction= False

thermal.bndCondIsTemperature=[1,1,0,0,0,0]
thermal.thermalBndCondValue=[0,0,0,0,0,0]
thermal.tsSafetyFactor=0
thermal.particleDensity=2600
thermal.particleT0=t0
thermal.particleCp=heatCap
thermal.particleK=thermalCond
thermal.particleAlpha =11.6e-3
thermal.useKernMethod=False

timing.reset()
#ThermalEngine.dead=0

flow.updateTriangulation=True
O.dt=1.
O.dynDt=False

flow.emulateAction()
flow.dead=1

#triax.goal2=-11000

def bodyByPos(x,y,z):
	cBody = O.bodies[1]
	cDist = Vector3(100,100,100)
	for b in O.bodies:
		if isinstance(b.shape, Sphere):
			dist = b.state.pos - Vector3(x,y,z)
			if np.linalg.norm(dist) < np.linalg.norm(cDist):
				cDist = dist
				cBody = b
	print('found closest body ', cBody.id, ' at ', cBody.state.pos)
	return cBody

# solution to the heat equation for constant initial condition , BCs=0, and using series for approx
def analyticalHeatSolution(x,t,u0,L,k):
	ns = np.linspace(1,1000,1000)
	solution = 0
	for i,n in enumerate(ns):
		integral = (-2./L)*u0*L*(np.cos(n*np.pi)-1.) / (n*np.pi)
		solution += integral * np.sin(n*np.pi*x/L)*np.exp((-k*(n*np.pi/L)**2)*t)
	return solution
  
# find 10 bodies along x axis
axis = np.linspace(mn[0], mx[0], num=11)
axisBodies = [None] * len(axis)
axisTrue = np.zeros(len(axis))
for i,x in enumerate(axis):
	axisBodies[i] = bodyByPos(x, mx[1]/2, mx[2]/2)
	axisTrue[i] = axisBodies[i].state.pos[0]
np.savetxt('txt'+timeStr+identifier+'/xdata.txt',axisTrue)
print("Axis length used for analy ", max(axisTrue)-min(axisTrue))
from yade import plot

## a function saving variables
def history():
	plot.addData(
		t=O.time,
		i = O.iter,
		temp1 = axisBodies[0].state.temp,
		temp2 = axisBodies[1].state.temp,
		temp3 = axisBodies[2].state.temp,
		temp4 = axisBodies[3].state.temp,
		temp5 = axisBodies[4].state.temp,
		temp6 = axisBodies[5].state.temp,
		temp7 = axisBodies[6].state.temp,
		temp8 = axisBodies[7].state.temp,
		temp9 = axisBodies[8].state.temp,
		temp10 = axisBodies[9].state.temp,
		temp11 = axisBodies[10].state.temp,
		AnalyTemp1 = analyticalHeatSolution(0,O.time,t0,mx[0],thermalDiff),
		AnalyTemp2 = analyticalHeatSolution(axisBodies[1].state.pos[0],O.time,t0,mx[0],thermalDiff),
		AnalyTemp3 = analyticalHeatSolution(axisBodies[2].state.pos[0]-min(axisTrue),O.time,t0,max(axisTrue)-min(axisTrue),thermalDiff),
		AnalyTemp4 = analyticalHeatSolution(axisBodies[3].state.pos[0]-min(axisTrue),O.time,t0,max(axisTrue)-min(axisTrue),thermalDiff),
		#bndFlux1 = thermal.thermalBndFlux[0],
		#bndFlux2 = thermal.thermalBndFlux[1]
		AnalyTemp5 = analyticalHeatSolution(mx[0],O.time,t0,mx[0],thermalDiff)
)
	plot.saveDataTxt('txt'+timeStr+identifier+'/conductionAnalyticalComparison.txt',vars=('t','i','temp1','temp2','temp3','temp4','temp5','temp6','temp7','temp8','temp9','temp10','temp11'))
  	#plot.addData(e22=-triax.strain[1],t=O.time,s22=-triax.stress(2)[1],p=flow.MeasurePorePressure((0.5,0.5,0.5)))

O.engines=O.engines+[PyRunner(iterPeriod=500,command='history()',label='recorder')]
##make nice animations:
#O.engines=O.engines+[PyRunner(iterPeriod=200,command='flow.saveVtk(withBoundaries=False)')]
VTKrec.dead=0
from yade import plot

plot.plots={'t':(('temp4','k-'),('temp3','r-'),('AnalyTemp4','k--'),('AnalyTemp3','r--'))} #
#plot.plots={'t':(('bndFlux1','k-'),('bndFlux2','r-'))} #
plot.plot()
O.saveTmp()
O.timingEnabled=1
from yade import timing
print("starting oedometer simulation")
O.run(200,1)
timing.stats()
