# -*- encoding=utf-8 -*-
from __future__ import print_function
######################################################################
# Compress the loose packing into dense packing. Each agglomerate is
# considered as clump in this stage
######################################################################
from builtins import zip
from builtins import range
from yade import export, ymport
import random
random.seed(1)

# add walls first
dim = (15, 15, 15)
walls = aabbWalls(((0, 0, 0), (dim)))
wallIds = O.bodies.append(walls)

# load spheres from file, including information of their agglomerates ids
attrs = []
print("\nPlease see README if you have errors when running this example.\n")
sp = ymport.textExt('/tmp/divided.txt', format='x_y_z_r_attrs', attrs=attrs)
n = max(int(a[0]) for a in attrs) + 1
colors = [randomColor() for _ in range(n)]
agglomerates = [[] for _ in range(n)]
for s, a in zip(sp, attrs):
	aa = int(a[0])
	s.agglomerate = aa
	s.shape.color = colors[aa]
	agglomerates[aa].append(s)
for a in agglomerates:
	O.bodies.appendClumped(a)

O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(), Bo1_Box_Aabb()]),
        InteractionLoop([Ig2_Sphere_Sphere_ScGeom(), Ig2_Box_Sphere_ScGeom()], [Ip2_FrictMat_FrictMat_FrictPhys()], [Law2_ScGeom_FrictPhys_CundallStrack()]),
        TriaxialStressController(
                thickness=0,
                stressMask=7,
                internalCompaction=False,
                label='compressor',
        ),
        NewtonIntegrator(damping=.6),
]
O.dt = PWaveTimeStep()

# compress the sample
compressor.goal1 = compressor.goal2 = compressor.goal3 = -1e-7
O.run(50000, True)
compressor.goal1 = compressor.goal2 = compressor.goal3 = -1e-5
O.run(30000, True)

# save the result, including information of agglomerates which the particle belongs to
export.textExt('/tmp/compressed.txt', 'x_y_z_r_attrs', attrs=['b.agglomerate'])

try:
	from yade import qt
	qt.View()
except:
	pass
