# -*- coding: utf-8 -*-

from __future__ import print_function
from yade import pack, utils, ymport, export

packing='parallellepiped_10'
DFN='persistentPlane30Deg'

############################ material definition
facetMat = O.materials.append(JCFpmMat(type=0,young=1,frictionAngle=radians(1),poisson=0.4,density=1))
def sphereMat(): return JCFpmMat(type=1,young=1,frictionAngle=radians(1),density=1,poisson=1,tensileStrength=1e6,cohesion=1e6,jointNormalStiffness=1,jointShearStiffness=1,jointTensileStrength=1e6,jointCohesion=1e6,jointFrictionAngle=1)

############################ Import of the sphere assembly
print("\n  In case of errors please look at README about generating parallellepiped_10.spheres file\n")
O.bodies.append(ymport.text(packing+'.spheres',scale=1,shift=Vector3(0,0,0),material=sphereMat)) #(-3,-4,-8)

#### some preprocessing (not mandatory)
dim=utils.aabbExtrema()
xinf=dim[0][0]
xsup=dim[1][0]
yinf=dim[0][1]
ysup=dim[1][1]
zinf=dim[0][2]
zsup=dim[1][2]

R=0
Rmax=0
numSpheres=0.
for o in O.bodies:
 if isinstance(o.shape,Sphere):
   o.shape.color=(0,0,1)
   numSpheres+=1
   R+=o.shape.radius
   if o.shape.radius>Rmax:
     Rmax=o.shape.radius
 else :
   o.shape.color=(0,0,0)
Rmean=R/numSpheres

print('number of spheres=', numSpheres, ' | Rmean=', Rmean, ' | dim=', dim)

############################ import stl file
O.bodies.append(ymport.stl(DFN+'.stl',color=(0.9,0.9,0.9),wire=False,material=facetMat)) 

############################ engines definition
interactionRadius=1.;
O.engines=[

	ForceResetter(),
	InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=interactionRadius,label='is2aabb'),Bo1_Facet_Aabb()]),
	InteractionLoop(
		[Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=interactionRadius,label='ss2d3dg'),Ig2_Facet_Sphere_ScGeom()],
		[Ip2_JCFpmMat_JCFpmMat_JCFpmPhys(cohesiveTresholdIteration=1,label='interactionPhys')],
		[Law2_ScGeom_JCFpmPhys_JointedCohesiveFrictionalPM(smoothJoint=True,label='interactionLaw')]
	),
	NewtonIntegrator(damping=1)

]

############################ timestep + opening yade windows
O.dt=0.001*utils.PWaveTimeStep()

from yade import qt
v=qt.Controller()
v=qt.View()

############################ Identification spheres on joint
#### color set for particles on joint
jointcolor1=(0,1,0)
jointcolor2=(1,0,0)
jointcolor3=(0,0,1)
jointcolor4=(1,1,1)
jointcolor5=(0,0,0)

#### first step-> find spheres on facet
O.step();

for i in O.interactions:
    ##if not i.isReal : continue
    if isinstance(O.bodies[i.id1].shape,Facet) and isinstance(O.bodies[i.id2].shape,Sphere): 
        vertices=O.bodies[i.id1].shape.vertices
        normalRef=vertices[0].cross(vertices[1]) # defines the normal to the facet normalRef
        nRef=normalRef/(normalRef.norm()) ## normalizes normalRef
        normalFacetSphere=i.geom.normal # geom.normal is oriented from id1 to id2 -> normalFacetSphere from facet (i.id1) to sphere (i.id2)

        if O.bodies[i.id2].state.onJoint==False : ## particles has not yet been identified as belonging to a joint plane
            O.bodies[i.id2].state.onJoint=True
            O.bodies[i.id2].state.joint=1
            O.bodies[i.id2].shape.color=jointcolor1
            if nRef.dot(normalFacetSphere)>=0 :
                O.bodies[i.id2].state.jointNormal1=nRef
            elif nRef.dot(normalFacetSphere)<0 :
                O.bodies[i.id2].state.jointNormal1=-nRef
        elif O.bodies[i.id2].state.onJoint==True :  ## particles has already been identified as belonging to, at least, 1 facet
            if O.bodies[i.id2].state.joint==1 and ((O.bodies[i.id2].state.jointNormal1.cross(nRef)).norm()>0.05) : ## particles has already been identified as belonging to only 1 facet
                O.bodies[i.id2].state.joint=2
                O.bodies[i.id2].shape.color=jointcolor2
                if nRef.dot(normalFacetSphere)>=0 :
                    O.bodies[i.id2].state.jointNormal2=nRef
                elif nRef.dot(normalFacetSphere)<0 :
                    O.bodies[i.id2].state.jointNormal2=-nRef
            elif O.bodies[i.id2].state.joint==2 and ((O.bodies[i.id2].state.jointNormal1.cross(nRef)).norm()>0.05) and ((O.bodies[i.id2].state.jointNormal2.cross(nRef)).norm()>0.05): ## particles has already been identified as belonging to more than 1 facet
                O.bodies[i.id2].state.joint=3
                O.bodies[i.id2].shape.color=jointcolor3
                if nRef.dot(normalFacetSphere)>=0 :
                    O.bodies[i.id2].state.jointNormal3=nRef
                elif nRef.dot(normalFacetSphere)<0 :
                    O.bodies[i.id2].state.jointNormal3=-nRef
            elif O.bodies[i.id2].state.joint==3 and ((O.bodies[i.id2].state.jointNormal1.cross(nRef)).norm()>0.05) and ((O.bodies[i.id2].state.jointNormal2.cross(nRef)).norm()>0.05) and ((O.bodies[i.id2].state.jointNormal3.cross(nRef)).norm()>0.05):
                O.bodies[i.id2].state.joint=4
                O.bodies[i.id2].shape.color=jointcolor5

#### second step -> find spheres interacting with spheres on facet (could be executed in the same timestep as step 1?)
for j in O.interactions:
    #if not i.isReal : continue
    if isinstance(O.bodies[j.id1].shape,Facet) and isinstance(O.bodies[j.id2].shape,Sphere):
        vertices=O.bodies[j.id1].shape.vertices
        normalRef=vertices[0].cross(vertices[1]) # defines the normal to the facet normalRef
        nRef=normalRef/(normalRef.norm()) ## normalizes normalRef
        if ((O.bodies[j.id2].state.jointNormal1.cross(nRef)).norm()<0.05) :
            jointNormalRef=O.bodies[j.id2].state.jointNormal1
        elif ((O.bodies[j.id2].state.jointNormal2.cross(nRef)).norm()<0.05) :
            jointNormalRef=O.bodies[j.id2].state.jointNormal2
        elif ((O.bodies[j.id2].state.jointNormal3.cross(nRef)).norm()<0.05) :
            jointNormalRef=O.bodies[j.id2].state.jointNormal3
        else : continue
        facetCenter=O.bodies[j.id1].state.pos
        #### seek for each sphere interacting with the identified sphere j.id2
        for n in O.interactions.withBody(j.id2) :
            if n.id1==j.id2 and isinstance(O.bodies[n.id2].shape,Sphere): 
                facetSphereDir=(O.bodies[n.id2].state.pos-facetCenter)
                if O.bodies[n.id2].state.onJoint==True :
                    if O.bodies[n.id2].state.joint==3 and ((O.bodies[n.id2].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id2].state.jointNormal2.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id2].state.jointNormal3.cross(jointNormalRef)).norm()>0.05):
                        O.bodies[n.id2].state.joint=4
                        O.bodies[n.id2].shape.color=jointcolor5
                    elif O.bodies[n.id2].state.joint==2 and ((O.bodies[n.id2].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id2].state.jointNormal2.cross(jointNormalRef)).norm()>0.05):
                        O.bodies[n.id2].state.joint=3
                        if facetSphereDir.dot(jointNormalRef)>=0:
                            O.bodies[n.id2].state.jointNormal3=jointNormalRef
                        elif facetSphereDir.dot(jointNormalRef)<0:
                            O.bodies[n.id2].state.jointNormal3=-jointNormalRef 
                    elif O.bodies[n.id2].state.joint==1 and ((O.bodies[n.id2].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) :
                        O.bodies[n.id2].state.joint=2
                        if facetSphereDir.dot(jointNormalRef)>=0:
                            O.bodies[n.id2].state.jointNormal2=jointNormalRef
                        elif facetSphereDir.dot(jointNormalRef)<0:
                            O.bodies[n.id2].state.jointNormal2=-jointNormalRef
                elif  O.bodies[n.id2].state.onJoint==False :
                    O.bodies[n.id2].state.onJoint=True
                    O.bodies[n.id2].state.joint=1
                    O.bodies[n.id2].shape.color=jointcolor4
                    if facetSphereDir.dot(jointNormalRef)>=0:
                        O.bodies[n.id2].state.jointNormal1=jointNormalRef
                    elif facetSphereDir.dot(jointNormalRef)<0:
                        O.bodies[n.id2].state.jointNormal1=-jointNormalRef

            elif n.id2==j.id2 and isinstance(O.bodies[n.id1].shape,Sphere): 
                facetSphereDir=(O.bodies[n.id1].state.pos-facetCenter)
                if O.bodies[n.id1].state.onJoint==True :
                    if O.bodies[n.id1].state.joint==3 and ((O.bodies[n.id1].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id1].state.jointNormal2.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id1].state.jointNormal3.cross(jointNormalRef)).norm()>0.05):
                        O.bodies[n.id1].state.joint=4
                        O.bodies[n.id1].shape.color=jointcolor5
                    elif O.bodies[n.id1].state.joint==2 and ((O.bodies[n.id1].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) and ((O.bodies[n.id1].state.jointNormal2.cross(jointNormalRef)).norm()>0.05):
                        O.bodies[n.id1].state.joint=3
                        if facetSphereDir.dot(jointNormalRef)>=0:
                            O.bodies[n.id1].state.jointNormal3=jointNormalRef
                        elif facetSphereDir.dot(jointNormalRef)<0:
                            O.bodies[n.id1].state.jointNormal3=-jointNormalRef
                    elif O.bodies[n.id1].state.joint==1 and ((O.bodies[n.id1].state.jointNormal1.cross(jointNormalRef)).norm()>0.05) :
                        O.bodies[n.id1].state.joint=2
                        if facetSphereDir.dot(jointNormalRef)>=0:
                            O.bodies[n.id1].state.jointNormal2=jointNormalRef
                        elif facetSphereDir.dot(jointNormalRef)<0:
                            O.bodies[n.id1].state.jointNormal2=-jointNormalRef
                elif  O.bodies[n.id1].state.onJoint==False :
                    O.bodies[n.id1].state.onJoint=True
                    O.bodies[n.id1].state.joint=1
                    O.bodies[n.id1].shape.color=jointcolor4
                    if facetSphereDir.dot(jointNormalRef)>=0:
                        O.bodies[n.id1].state.jointNormal1=jointNormalRef
                    elif facetSphereDir.dot(jointNormalRef)<0:
                        O.bodies[n.id1].state.jointNormal1=-jointNormalRef


#### for visualization:
#bj=0
#vert=(0.,1.,0.)
#hor=(0.,1.,1.)
#for o in O.bodies:
    #if o.state.onJoint==True : # or o.shape.name=='Facet':
	##if  o.shape.name=='Facet':
	    ##o.shape.wire=True
	##o.state.pos+=(0,50,0)
	##bj+=1
	#if o.state.jointNormal1.dot(hor)>0 :
	    ##o.state.pos+=(0,50,0)
	    #o.shape.color=jointcolor1
	#elif o.state.jointNormal1.dot(hor)<0 :
	    ##o.state.pos+=(0,55,0)
	    #o.shape.color=jointcolor2
	#if o.mat.type>2 :
	    #bj+=1
	    #o.shape.color=jointcolor5
	    ##print o.state.jointNormal.dot(vert)


#### Save text file with informations on each sphere
export.text(packing+'_'+DFN+'.spheres')
export.textExt(packing+'_'+DFN+'_jointedPM.spheres',format='jointedPM')

O.wait()
