#Import Libraries

from yade import pack,plot
import math
import random as rand
import numpy as np
from scipy.stats import gaussian_kde

# Time step
O.dt=1e-6

# Skeleton creation
title='FILETAIEAIE'
folder=os.path.dirname(sys.argv[0])+"Results/"+str(title)+'/'
if os.path.exists(folder)==False:
    os.mkdir(folder)
    os.mkdir(folder+'VTK/')

##############################
## PARAMETERS OF SIMULATION ##
##############################
###Configuration : inclined channel
slope = math.radians(31)    #Slope of the channe, in radian
length = 90 #Length of the channe, in m
width = 2   #Width of the channel, in m
height = 2  #Height of the channel, in m
gravityVector = Vector3(9.81*sin(slope),0.0,-9.81*cos(slope))   #Gravity vector to consider a channel inclined with slope angle 'slope'

### Materials
# Channel
chanFrictAng =  70.

#Fluid
fluidDensity =  1200.   #Density of the fluid, in kg/m3, used for the computation of the dragforce
m_viscosity =   2   #Test avec 0.022

# Debris
diameterPart =  0.1 #Diameter of the particles, in m
m_roughness =   0.01    #Fraction on the radius
sphDensity =    2500
sphYoung =  1e7
sphPoisson =    0.2
sphFrictAng =   35.

O.materials.append(
        FrictMat(
            density=sphDensity,
            frictionAngle=math.radians(sphFrictAng),
            poisson=sphPoisson,
            young=sphYoung,
            label='Frict'
        ))

O.materials.append(
                FrictMat(
                        density=2500,
                        frictionAngle=math.radians(chanFrictAng),
                        poisson=0.2,
                        young=1e7,
                        label='FrictBase'
                ))

O.materials.append(
                FrictMat(
                        density=2500,
                        frictionAngle=math.radians(0.),
                        poisson=0.2,
                        young=1e7,
                        label='FrictWall'
                ))

O.materials.append(
                FrictMat(
                        density=2500,
                        frictionAngle=math.radians(20.),
                        poisson=0.2,
                        young=1e7,
                        label='FrictRigid'
                ))

########################
## FRAMEWORK CREATION ##
########################

# Channel creation
boxColor = (104/255., 111/255., 140/255.)
chanBack = O.bodies.append(aabbWalls([(0,0,0),(0,width,height)],thickness=0.0,material='FrictWall',oversizeFactor=1.0,color=boxColor))
chanS1 = O.bodies.append(aabbWalls([(0,0,0),(length,0,height)],thickness=0.0,material='FrictWall',oversizeFactor=1.0,color=boxColor))
chanB = O.bodies.append(aabbWalls([(0,0,0),(length,width,0)],thickness=0.0,material='FrictBase',oversizeFactor=1.0,color=boxColor))
chanS2 = O.bodies.append(aabbWalls([(0,width,0),(length,width,height)],thickness=0.0,material='FrictWall',oversizeFactor=1.0,color=boxColor))
chanBack = O.bodies.append(aabbWalls([(length,0,0),(length,width,height)],thickness=0.0,material='FrictWall',oversizeFactor=1.0,color=boxColor))
# Deposit box
DepBox = O.bodies.append(aabbWalls([(0.,0.,-0.1),(0.1*length,width,2*height)],thickness=0.0,material='FrictWall',oversizeFactor=1.0,color=boxColor))

# Rigid wall
rigWall = O.bodies.append(
        aabbWalls(
        [(0.9*length,0,0),(0.9*length+0.1,width,height/2)],
        thickness=0.,
        material='FrictRigid',
        oversizeFactor=1.0,
        color=boxColor))



# Creation of sphere pack
partCloud = pack.SpherePack()
partNumber = 10000
psdSizes = [0.08,0.16,0.32,0.64]#[0.00125,0.0025,0.005,0.01,0.02,0.04,0.08,0.16,0.32] #couper au centimetre
psdCumm = [0,0.53,0.86,1.  ]#[  0.01,  0.05, 0.17, 0.26,0.34,0.43,0.53,0.86,1.  ]
#partCloud.makeCloud(minCorner=(0.,0.,0.),maxCorner=(length/10.,width,2*height), psdSizes=psdSizes, psdCumm=psdCumm, distributeMass=True,seed=23)
partCloud.makeCloud(minCorner=(length/10./3.*2.,0.,0.),maxCorner=(length/10.,width,height), rMean=diameterPart/2.0)
spIds = partCloud.toSimulation(material='Frict',color=[205./255,175./255,149./255])
allSp = spIds
#####################
## SIMULATION LOOP ##
#####################

# Definition of simulation's variable
Factor = 1.5    #Factor ratio for the distance detection of distant interactions
theta_method = 0.55 #1 : Euler Backward, 0.5 : trapezoidal rule, 0 : not used, 0.55 : suggested optimum
resolution = 1      #Change normal component resolution method, 0: Iterative exact resolution (theta method, linear contact), 1: Newton-Rafson dimentionless resolution (theta method, linear contact), 2: Newton-Rafson with nonlinear surface deflection (Hertzian-like contact)
NewtonRafsonTol = 1e-6  #Precision of the resolution of Newton-Rafson's method
vel_front = Vector3(0.,0.,0.)

# Engine loop
O.engines=[
    ForceResetter(),
    InsertionSortCollider([
        Bo1_Sphere_Aabb(aabbEnlargeFactor=Factor, label="aabb"),
        Bo1_Box_Aabb(),
        Bo1_Wall_Aabb()],
        verletDist=-0.1,
        allowBiggerThanPeriod=False),

    InteractionLoop(
    [
        Ig2_Sphere_Sphere_ScGeom6D(interactionDetectionFactor=Factor,label="Ig2"),
        Ig2_Box_Sphere_ScGeom6D()
    ],
        [
        Ip2_FrictMat_FrictMat_LubricationPhys(eta=m_viscosity,eps=m_roughness,label="Phys_Lub"),
    ],
        [
        Law2_ScGeom_ImplicitLubricationPhys(
                        activateNormalLubrication=True,
                        activateTangencialLubrication=True,
                        activateTwistLubrication=True,
                        activateRollLubrication=True,
                        theta=theta_method,
                        resolution=resolution,
                        warnedOnce=True,
                        maxSubSteps=20,
                        NewtonRafsonTol=NewtonRafsonTol,
                        debug=False)
    ]),

    #DragEngine(Rho=1200,Cd=0.47,V=vel_front, label='Drag', dead=True),
    NewtonIntegrator(gravity=gravityVector, damping = 0.4, label='newton'),

    # Launching Python functions inside the loop
    PyRunner(command='gravityDeposit()',iterPeriod=500,label='gravDep',dead=False),
    PyRunner(command='Fx, Fy, Fz=Force_wall()',iterPeriod=1000,label='FWall',dead=True),
    PyRunner(command='deadZone()',iterPeriod=1000,label='deadZone_engine',dead=True),
    PyRunner(command='max_pos_x,height_front, vel_front=front()',iterPeriod=1000,label='dista',dead=True),
    PyRunner(command='vel_profil()',iterPeriod=1000,label='velProfil',dead=True),
    PyRunner(command='shearLub_tot,shearCon_tot=bedshear_total()',iterPeriod=1000,label='bedshear_tot_engine',dead=True),
    PyRunner(command='shearLub1,shearCon1=bedshear_at_x(0.15*length)',iterPeriod=1000,label='bedshear_1_engine',dead=True),
    PyRunner(command='shearLub2,shearCon2=bedshear_at_x(30)',iterPeriod=1000,label='bedshear_2_engine',dead=True),
    PyRunner(command='shearLub3,shearCon3=bedshear_at_x(50)',iterPeriod=1000,label='bedshear_3_engine',dead=True),
    PyRunner(command='shearLub4,shearCon4=bedshear_at_x(75)',iterPeriod=1000,label='bedshear_4_engine',dead=True),
    VTKRecorder(fileName=folder+'/VTK/vtk-',virtPeriod=0.1),
#   PyRunner(command='Epot,Ekin,Etot=Energy()',iterPeriod=500),
#   PyRunner(command='addPlotData()',iterPeriod=1000), #ne pas oublier de commenter
    ]

# Caracteristics of the drag engine
#Drag.ids=spIds
#Drag.Rho=fluidDensity
#Drag.Cd=0.47

#########################
## FUNCTION DEFINITION ##
#########################

def gravityDeposit():
    if O.iter>5000 and kineticEnergy()<10 and unbalancedForce()<1:
        #O.pause()
        for i in DepBox:
            O.bodies.erase(i)
        gravDep.dead=True
        #Drag.dead=False
        dista.dead=False
        velProfil.dead=False
        bedshear_tot_engine.dead=False
        bedshear_1_engine.dead=False
        bedshear_2_engine.dead=False
        bedshear_3_engine.dead=False
        bedshear_4_engine.dead=False
        newton.damping=0.
        FWall.dead=False
        Phys_Lub.eta=m_viscosity
#       O.save("deposit-mono5cm.yade")
        print "Gravity deposit finished & gate opened"

def addPack(minX,minY,minZ,maxX,maxY,maxZ,num):
    global spIds
    PackAdd = pack.SpherePack()
    PackAdd.makeCloud(minCorner=(minX,minY,minZ),maxCorner=(maxX,maxY,maxZ),rMean=diameterPart/2.,num=num)
    spIdsAdd=PackAdd.toSimulation(material=mat,color=[205./255,175./255,149./255])
    spIds=spIds+spIdsAdd
    #Drag.ids=spIds

xs=np.linspace(0,length,int(length)*10)
zs=np.linspace(0,height,int(height)*100)
max_pos_x, height_front, vel_front, max_vel_front=0.,0.,Vector3(0.,0.,0.),0.
def front():
    global max_vel_front
    fvel=open(folder+"front.txt","a")
    fvel3D=open(folder+"front3D.txt","a")
#   for i in spIds:
#       O.bodies[i].shape.color=[205./255,175./255,149./255]
    list_x,list_z=[],[]
    for i in spIds:
        list_x.append(O.bodies[i].state.pos[0])
        list_z.append(O.bodies[i].state.pos[2])
    density_x = gaussian_kde(list_x)
    density_x.covariance_factor = lambda : .05
    density_x._compute_covariance()
    density_z = gaussian_kde(list_z)
    density_z.covariance_factor = lambda : .05
    density_z._compute_covariance()

    max_density_x=max(density_x(xs))
    max_pos_x=xs[max([i for i,x in enumerate(density_x(xs)) if x>0.5*max_density_x])]

    height=max([x for x in zs if 0.95>density_z.integrate_box_1d(0,x)])+diameterPart/2.

    vel_front,a=Vector3(0.,0.,0.),0
    for i in spIds:
        if abs(O.bodies[i].state.pos[0] - max_pos_x) <=0.5:
#           O.bodies[i].shape.color=[0./255,0./255,255./255]
            vel_front+=O.bodies[i].state.vel
            a+=1
    if a!=0:
        vel_front=vel_front/float(a)
        #Drag.V=vel_front

    max_vel_front=max(max_vel_front,vel_front[0])

    fvel.write(str(O.time)+";"+str(max_pos_x)+";"+str(height)+";"+str(vel_front[0])+"\n")
    fvel3D.write(str(O.time)+";"+str(vel_front[0])+";"+str(vel_front[1])+";"+str(vel_front[2])+"\n")
    fvel.close()
    fvel3D.close()

    return max_pos_x,height, vel_front

def distance(a,b):
    return sqrt((a.state.pos[0]-b.state.pos[0])**2+(a.state.pos[1]-b.state.pos[1])**2+(a.state.pos[2]-b.state.pos[2])**2)

deadZoneIds=[]
def deadZone():
    global spIds
    for i in spIds:
        if (O.bodies[i].state.vel[0]/float(max_vel_front)) < 0.05:
            spIds.remove(i)
            deadZoneIds.append(i)
    if spIds==[]:
        dista.dead=True
        velProfil.dead=True
    #Drag.ids=spIds

def vel_profil():
    velocity_profil=open(folder+"velocity_profil.txt","a") #ecriture dans un fichier velocity_profil.txt
    velocity_profil.write(str(O.time)+";"+str(max_pos_x)+";")
    l=[]
    for i in spIds:
        if max_pos_x-1.2 <= O.bodies[i].state.pos[0] <= max_pos_x-0.8:
            l.append(i)

    list_vel_x,list_vel_y,list_vel_z="0","0","0"
    for h in range(int(height_front/diameterPart)):
        vel_x,vel_y,vel_z,a=0.,0.,0.,0
        for i in l:
            if h*diameterPart <= O.bodies[i].state.pos[2] <= (h+1)*diameterPart:
                vel_x+=O.bodies[i].state.vel[0]
                vel_y+=O.bodies[i].state.vel[1]
                vel_z+=O.bodies[i].state.vel[2]
                a+=1
        if a!=0:
            list_vel_x+=":"+str(vel_x/a)
            list_vel_y+=":"+str(vel_y/a)
            list_vel_z+=":"+str(vel_z/a)
    velocity_profil.write(str(height_front)+";"+str(list_vel_x)+";"+str(list_vel_y)+";"+str(list_vel_z)+"\n")
    velocity_profil.close()

FtotX,FtotY,FtotZ,tF=[],[],[],[]
def Force_wall():
    f=open(folder+"Force_wall.txt","a")
    ### Calculate each component of the total NORMAL force between impacting particles and the wall:
    ### Diviser le mur en multiple tranche

    Fx=0
    Fy=0
    Fz=0
    for i in O.interactions:
        if i.id1 == rigWall[0]:
            deadZone_engine.dead=False
            Fx+= i.phys.normalForce[0]
            Fy+= i.phys.normalForce[1]
            Fz+= i.phys.normalForce[2]
    FtotX.append(Fx)
    FtotY.append(Fy)
    FtotZ.append(Fz)
    tF.append(O.time)
    f.write(str(O.time)+";"+str(Fx)+";"+str(Fy)+";"+str(Fz)+"\n")
    f.close()
    return Fx, Fy, Fz

shearLub_tot,shearCon_tot=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
def bedshear_total():
    shearLub,shearCon=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
    fbedshear=open(folder+"bedshear.txt","a")
    for i in O.interactions:
        if i.id1 in chanB:
            shearLub+=i.phys.shearLubricationForce
            shearCon+=i.phys.shearContactForce
    fbedshear.write(str(O.time)+";"+str(shearLub[0])+";"+str(shearLub[1])+";"+str(shearLub[2])+";"+str(shearCon[0])+";"+str(shearCon[1])+";"+str(shearCon[2])+"\n")
    fbedshear.close()
    return shearLub,shearCon

shearLub1,shearCon1=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
shearLub2,shearCon2=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
shearLub3,shearCon3=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
shearLub4,shearCon4=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
shearLub5,shearCon5=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
def bedshear_at_x(x):
    fbedshear=open(folder+"bedshear_at_"+str(x)+".txt","a")
    shearLub,shearCon=Vector3(0.,0.,0.),Vector3(0.,0.,0.)
    for i in O.interactions:
        if (i.id1 in chanB) and (x-0.5 <= O.bodies[i.id2].state.pos[0] <= x+0.5):
            shearLub+=i.phys.shearLubricationForce
            shearCon+=i.phys.shearContactForce
    fbedshear.write(str(O.time)+";"+str(shearLub[0])+";"+str(shearLub[1])+";"+str(shearLub[2])+";"+str(shearCon[0])+";"+str(shearCon[1])+";"+str(shearCon[2])+"\n")
    fbedshear.close()
    return shearLub,shearCon

Epot,Ekin,Etot=0.,0.,0.
def Energy():
    Epot,Ekin,Etot=0.,0.,0.
    for i in spIds:
        p=Vector3(length-O.bodies[i].state.pos[0],O.bodies[i].state.pos[1],O.bodies[i].state.pos[2])
        Epot+=O.bodies[i].state.mass*Vector3.dot(gravityVector,p)
        Ekin+=0.5*O.bodies[i].state.mass*O.bodies[i].state.vel.norm()**2
    return Epot,Ekin,Epot+Ekin

def DragForce(rho, Cd):
    for i in spIds:
        relVel = -(Vector3(vel_front,0.,0.)-O.bodies[i].state.vel)
        drag = -0.5 * relVel.squaredNorm() * relVel.normalized() * rho * Cd
        O.forces.addF(i,drag)



def addPlotData():
#       plot.addData(i=O.iter, unbalanced=unbalancedForce(),t=O.time,front_vel=front_vel,tt=O.time,shearLub_tot=shearLub_tot[0],shearCon_tot=shearCon_tot[0],ttt=O.time,shearLub1=shearLub1[0],shearCon1=shearCon1[0],tttt=O.time,Epot=Epot,Ekin=Ekin,Etot=Etot, **O.energy)
    plot.addData(i=O.iter, unbalanced=unbalancedForce(),t=O.time,vel_front=vel_front[0],tt=O.time,shearLub_tot=shearLub_tot[0],shearCon_tot=shearCon_tot[0],ttt=O.time,shearLub1=shearLub1[0],shearCon1=shearCon1[0],tttt=O.time,Epot=Epot,Ekin=Ekin,Etot=Etot, **O.energy)



O.trackEnergy=True

#plot.plots={'i':('unbalanced'),'tttt':('Epot',None,('Etot','g')),'tt':'vel_front'}
#plot.labels={'vel_front':'vel_front '+title}
#plot.plots={'t':'Epot','tt':'Ekin','ttt':'Etot'}
#plot.plot(subPlots=False)

f=open(folder+"param.txt","w")
f.write("###############################"+
    "\n## Parameters of the channel ##"+
    "\n###############################"+
    "\nlength = "+str(length)+" m"+
    "\nwidth = "+str(width)+" m"+
    "\nheight = "+str(height)+" m"+
    "\nslope = "+str(math.degrees(slope))+" degrees"+
        "\nchanFricAng = "+str(chanFrictAng)+" degrees"+
    "\n\n###############################"+
    "\n## Parameters of Lubrication ##"+
    "\n###############################"+
    "\nviscosity = "+str(m_viscosity)+" Pa.s"
    "\nroughness = "+str(m_roughness)+" % of diameterPart"
    "\n\n################################"+
    "\n## Parameters of the material ##"+
    "\n################################"+
    "\n\nType of material = Frict"+
    "\ndensity = "+str(sphDensity)+" kg/m^3"+
    "\ndiameter = "+str(diameterPart)+" m"+
    "\nyoung = "+str(sphYoung)+" Pa"+
    "\npoisson = "+str(sphPoisson)+
    "\nsphFricAng = "+str(sphFrictAng)+" degrees")
f.close()