# -*- encoding=utf-8 -*-
# 2021 J.Kozicki, A.Gladky, K.Thoeni "Implementation of high-precision computation capabilities into the dynamic simulation framework YADE" (to be submitted)
# The script is tested with yade commit 7f645d64f (April 2021)
from yade import plot, math

# initialize all floating point variables with Real(arg) to avoid precision loss
from yade.math import Real
import sys
import csv
import lzma

if(math == sys.modules['math']):
    raise RuntimeError(
        "To avoid precision problems it is better to avoid the python math (double only) module and use yade.math (faster) or mpmath (slower) instead.")

### set parameters ###
L = Real('0.1')      # length [m]
n = 4                # number of nodes for the length [-]
r = L/100            # radius [m]
g = Real('9.81')     # gravity
inclination = math.radiansHP1(20)  # Initial inclination of rods [degrees]
color = [1, 0.5, 0]  # Define a color for bodies

O.dt = Real('1e-05') # time step
damp = Real('1e-1')  # damping. It is interesting to examine  damp = 0

totalSecDuration = 100
totalIterDuration = int(totalSecDuration/O.dt)
iterPyRunnerPeriod = int(1e-4/O.dt)

O.engines = [  # define engines, main functions for simulation
    ForceResetter(),
    PyRunner(command='addPlotData()', initRun=True,
             iterPeriod=iterPyRunnerPeriod, nDo=1),  # run this at the beginning to get Eref=Epot(t=0)
    InsertionSortCollider([Bo1_Sphere_Aabb()],
                          label='ISCollider', avoidSelfInteractionMask=True),
    InteractionLoop(
        [Ig2_Sphere_Sphere_ScGeom6D()],
        [Ip2_CohFrictMat_CohFrictMat_CohFrictPhys(
            setCohesionNow=True, setCohesionOnNewContacts=False)],
        [Law2_ScGeom6D_CohFrictPhys_CohesionMoment()]
    ),
    NewtonIntegrator(gravity=(0, -g, 0), damping=damp, label='newton'),
    PyRunner(command='addPlotData()', initRun=True,
             iterPeriod=iterPyRunnerPeriod, firstIterRun=0),
    PyRunner(command='getPrevVel()', initRun=True,
             iterPeriod=iterPyRunnerPeriod, firstIterRun=iterPyRunnerPeriod-1),
    PyRunner(command='saveResults()', initRun=True,
             iterPeriod=iterPyRunnerPeriod, firstIterRun=totalIterDuration-1)
]

# define material:
O.materials.append(CohFrictMat(young=1e5, poisson=0, density=1e1,
                               frictionAngle=math.radiansHP1(0), normalCohesion=1e7,
                               shearCohesion=1e7, momentRotationLaw=False, label='mat'))

# create spheres
nodeIds = []
for i in range(0, n):
    nodeIds.append(O.bodies.append(sphere([i*L/n*math.cos(inclination),
        i*L/n*math.sin(inclination), 0], r, wire=False, fixed=False,
        material='mat', color=color)))

# create rods
for i, j in zip(nodeIds[:-1], nodeIds[1:]):
    inter = createInteraction(i, j)
    inter.phys.unp = -(O.bodies[j].state.pos-O.bodies[i].state.pos).norm() + O.bodies[i].shape.radius+O.bodies[j].shape.radius

# Set a fixed node
O.bodies[0].dynamic = False

# Calculate angle between three points: a, b, c
def calculateAngle(a, b, c):
    ba = a - b
    bc = c - b
    cosineAngle = ba.dot(bc) / (ba.norm() * bc.norm())

    if (cosineAngle > 1):
        angle = math.acos(1)
    elif (cosineAngle < -1):
        angle = math.acos(-1)
    else:
        angle = math.acos(cosineAngle)
    return (math.degrees(angle))

# Calculate potential energy
def calculatePotentialEnergy():
    potEnergy = Real(0)
    for i in range(0, n):
        h1 = O.bodies[i].state.pos[1]  # mass is lumped into nodes
        h2 = L/n*i
        potEnergy += O.bodies[i].state.mass * g * (h1 + h2)
    return potEnergy

# Calculate elastic energy
def calculateElasticEnergy():  # normal force only
    elEnergy = Real(0)
    for i in O.interactions:
        elEnergy += i.phys.normalForce.squaredNorm()/(2 * i.phys.kn)
    return elEnergy


dataToSave = []

iter_prev = 0
pos_prev = {'1': [0, 0],
            '2': [0, 0],
            '3': [0, 0]}
vel_prev = {'1': [0, 0],
            '2': [0, 0],
            '3': [0, 0]}
ang1_prev = 0
ang2_prev = 0

kin_prev = 0
pot_prev = 0
elast_prev = 0

# Get velocities for the previous step


def getPrevVel():
    global iter_prev
    global pos_prev
    global vel_prev
    global ang1_prev
    global ang2_prev
    global kin_prev
    global pot_prev
    global elast_prev
    iter_prev = O.iter
    pos_prev = {'1': [O.bodies[1].state.pos[0], O.bodies[1].state.pos[1]],
                '2': [O.bodies[2].state.pos[0], O.bodies[2].state.pos[1]],
                '3': [O.bodies[3].state.pos[0], O.bodies[3].state.pos[1]]}
    vel_prev = {'1': [O.bodies[1].state.vel[0], O.bodies[1].state.vel[1]],
                '2': [O.bodies[2].state.vel[0], O.bodies[2].state.vel[1]],
                '3': [O.bodies[3].state.vel[0], O.bodies[3].state.vel[1]]}
    ang1_prev = calculateAngle(O.bodies[0].state.pos,
                               O.bodies[1].state.pos, O.bodies[2].state.pos)
    ang2_prev = calculateAngle(O.bodies[1].state.pos,
                               O.bodies[2].state.pos, O.bodies[3].state.pos)
    kin_prev = kineticEnergy() # not used in the paper, only used for cross-checking
    pot_prev = calculatePotentialEnergy()
    elast_prev = calculateElasticEnergy()

# Function to accumulate and add data to plot


def addPlotData():
    ang1 = calculateAngle(O.bodies[0].state.pos,
                          O.bodies[1].state.pos, O.bodies[2].state.pos)
    ang2 = calculateAngle(O.bodies[1].state.pos,
                          O.bodies[2].state.pos, O.bodies[3].state.pos)

    curTime = O.iter*O.dt
    kin = kineticEnergy()
    pot = calculatePotentialEnergy()
    elast = calculateElasticEnergy()

    dataCurrent  = {'time': curTime, 'iter_prev': iter_prev, 'iter': O.iter,
                    'pos1_prev_x': pos_prev['1'][0], 'pos1_prev_y': pos_prev['1'][1],
                    'pos2_prev_x': pos_prev['2'][0], 'pos2_prev_y': pos_prev['2'][1],
                    'pos3_prev_x': pos_prev['3'][0], 'pos3_prev_y': pos_prev['3'][1],
                    'pos1_x': O.bodies[1].state.pos[0], 'pos1_y': O.bodies[1].state.pos[1],
                    'pos2_x': O.bodies[2].state.pos[0], 'pos2_y': O.bodies[2].state.pos[1],
                    'pos3_x': O.bodies[3].state.pos[0], 'pos3_y': O.bodies[3].state.pos[1],
                    'vel1_prev_x': vel_prev['1'][0],    'vel1_prev_y': vel_prev['1'][1],
                    'vel2_prev_x': vel_prev['2'][0],    'vel2_prev_y': vel_prev['2'][1],
                    'vel3_prev_x': vel_prev['3'][0],    'vel3_prev_y': vel_prev['3'][1],
                    'vel1_x': O.bodies[1].state.vel[0], 'vel1_y': O.bodies[1].state.vel[1],
                    'vel2_x': O.bodies[2].state.vel[0], 'vel2_y': O.bodies[2].state.vel[1],
                    'vel3_x': O.bodies[3].state.vel[0], 'vel3_y': O.bodies[3].state.vel[1],
                    'ang1_prev': ang1_prev,
                    'ang2_prev': ang2_prev,
                    'ang1': ang1,
                    'ang2': ang2,
                    'kin_prev': kin_prev,
                    'pot_prev': pot_prev,
                    'elast_prev': elast_prev,
                    'kin': kin,
                    'pot': pot,
                    'elast': elast
                    }
    dataToSave.append(dataCurrent)
    plot.addData(
        ang1=ang1, ang2=ang2, time1=curTime, time2=curTime, kin=kin, pot=pot, elast=elast
    )

def infoString():
    return '_digits_'+str(yade.math.getDigits10(1))+'_dt_'+str(O.dt)

def saveResults():
    fname1 = 'data_angles'+infoString()
    fname2 = 'data_energy'+infoString()+'.csv.xz'
    print('\n'+'-'*70)
    print(int(totalIterDuration), "iterations finished. Saving results to:\n",fname1,"\n",fname2)
    print('-'*70)
    plot.saveGnuplot(fname1)
    with lzma.open(fname2, mode='wt') as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=list(dataToSave[0].keys()))
        writer.writeheader()
        for i in dataToSave:
            writer.writerow(i)
    print('\n'+'-'*70)
    print('data saved')
    print('-'*70)

plot.plots = {'time1': ('ang1', 'ang2'), 'time2': ('kin', 'pot', 'elast')}
plot.plot()

O.saveTmp()

if(opts.nogui==False):
    from yade import qt
    qt.View()                    # create a GUI view
    Gl1_Sphere.stripes = True    # mark spheres with stripes
    rr = qt.Renderer()           # get instance of the renderer
    rr.intrAllWire = True        # draw wires
    rr.intrPhys = True           # draw the normal forces between the spheres.


print('-'*70)
print("Will now run",int(totalIterDuration),"iterations, which is",totalSecDuration,"seconds.")
print("Using time step = ",O.dt)
print('-'*70)

O.run(totalIterDuration)

