# -*- encoding=utf-8 -*-

from yade import pack, export, qt
import gts, os, random, itertools
from numpy import *

# Parameters
tc = 0.001  # collision time
en = .3  # normal restitution coefficient
es = .3  # tangential restitution coefficient
frictionAngle = radians(35)  #
density = 2700
# facets material
facetMat = O.materials.append(ViscElMat(frictionAngle=frictionAngle, tc=tc, en=en, et=es))
# default spheres material
dfltSpheresMat = O.materials.append(ViscElMat(density=density, frictionAngle=frictionAngle, tc=tc, en=en, et=es))

O.dt = .01 * tc  # time step

Rs = 0.1  # particle radius

# Create geometry
plnSurf = pack.sweptPolylines2gtsSurface([[Vector3(-.5, 0, 0), Vector3(.5, 0, 0), Vector3(.5, 0, -.5), Vector3(-.5, 0, -.5)]], capStart=True, capEnd=True)
plnIds = O.bodies.append(pack.gtsSurface2Facets(plnSurf, material=facetMat, color=(0, 1, 0)))

plnSurf1 = pack.sweptPolylines2gtsSurface(
        [[Vector3(-.5, -.5, -.5), Vector3(.5, -.5, -.5), Vector3(.5, 1.5, -.5),
          Vector3(-.5, 1.5, -.5)]], capStart=True, capEnd=True
)
plnIds1 = O.bodies.append(pack.gtsSurface2Facets(plnSurf1, material=facetMat, color=(0, 1, 0)))

# Create clumps
clpId, sphId = O.bodies.appendClumped([sphere(Vector3(0, Rs * 2 * i, Rs * 2), Rs, material=dfltSpheresMat) for i in range(4)])

# Create engines
O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(), Bo1_Facet_Aabb()]),
        InteractionLoop(
                [Ig2_Sphere_Sphere_ScGeom(), Ig2_Facet_Sphere_ScGeom()],
                [Ip2_ViscElMat_ViscElMat_ViscElPhys()],
                [Law2_ScGeom_ViscElPhys_Basic()],
        ),
        NewtonIntegrator(damping=0, gravity=[0, 0, -9.81]),
]

renderer = qt.Renderer()
qt.View()
O.saveTmp()

clump = O.bodies[clpId]
spheres = [O.bodies[id] for id in sphId]

print('\nClump:\n======')
#print('\nMaterial:')
if clump.mat:
	print(clump.mat.dict())
else:
	print('no material')

print('\nshape:')
if clump.shape:
	print(clump.shape.dict())
else:
	print('no shape')

print('\nState:')
if clump.state:
	print(clump.state.dict())
else:
	print('no state')

print('\nControl:')
mass = sum([s.state.mass for s in spheres])
xm_ = [s.state.pos[0] * s.state.mass for s in spheres]
ym_ = [s.state.pos[1] * s.state.mass for s in spheres]
zm_ = [s.state.pos[2] * s.state.mass for s in spheres]
centroid = Vector3(sum(xm_) / mass, sum(ym_) / mass, sum(zm_) / mass)


def sphereInertiaTensor(p, m, r, c):
	''' Inertia tensor sphere with position p, mass m and radus r relative point c '''
	I = zeros((3, 3))
	rp = array(p) - array(c)
	I[0, 0] = 2. / 5. * m * r * r + m * (rp[1]**2 + rp[2]**2)
	I[1, 1] = 2. / 5. * m * r * r + m * (rp[0]**2 + rp[2]**2)
	I[2, 2] = 2. / 5. * m * r * r + m * (rp[0]**2 + rp[1]**2)
	for k, l in itertools.product(arange(0, 3), arange(0, 3)):
		if not k == l:
			I[k, l] = -m * rp[k] * rp[l]
	return I


I = zeros((3, 3))
for s in spheres:
	I += sphereInertiaTensor(s.state.pos, s.state.mass, s.shape.radius, centroid)
I_eigenvalues, I_eigenvectors = linalg.eig(I)

print('mass = ', mass)
print('centroid = ', centroid)
print('I = \n', I)
print('I_eigenvalues = ', I_eigenvalues)
print('I_eigenvectors = \n', I_eigenvectors)
