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

readParamsFromTable(roughness=1.e-2*0, eta=100., a=1.e-3, rho=1e5, young=1.e9, shearRate=30)

from yade.params import table

from yade import plot
import math
import os

m_simuParams = {k: yade.params.table.__dict__[k] for k in table.__dict__['__all__']};
m_savefile = "data_shearTwoSpheres";

mul = 100;
dist = 0.05;

i = 0;
while os.path.exists(m_savefile+'.'+str(i)+'.txt'):
    i=i+1;

m_savefile += '.'+str(i)+'.txt';
open(m_savefile, 'a').close(); # Create the empty file asap

def SaveAndQuit():
    print("Quit condition reach.");
    saveToFile();
    O.stopAtIter = O.iter+1;

def upGraph():
        
        norm = 0;
        tangContact = Vector3(0,0,0);
        tangLub = Vector3(0,0,0);
        normContact = 0;
        normLub = 0;
        ue = 0;
        u = 0;
        
        [nc, sc, nl, sl] =  Law2_ScGeom_ImplicitLubricationPhys.getStressForEachBody();
        
        if O.interactions.has(1,0):
            if O.interactions[1,0].isReal:
                g = O.interactions[1,0].geom;
                ph = O.interactions[1,0].phys;

                norm = ph.normalForce.dot(g.normal);
                normContact = ph.normalContactForce.dot(g.normal);
                normLub = ph.normalLubricationForce.dot(g.normal);
                tangContact = ph.shearContactForce;
                tangLub = ph.shearLubricationForce;
                ue = ph.ue;
                u = ph.u;
        
        plot.addData(
                Fn = norm,
                ue = ue,
                Fnc = normContact,
                Fnl = normLub,
                Fsl = tangLub,
                Fsc = tangContact,
                nc = nc[0],
                nl = nl[0],
                u = u,
                t = O.time,
                t2 = O.time,
                t3 = O.time,
                t4 = O.time,
                pos0 = O.bodies[0].state.pos,
                pos1 = O.bodies[1].state.pos)

        if(O.bodies[1].state.pos[0] - O.bodies[0].state.pos[0] > 6.*table.a):
            SaveAndQuit();

       

def saveToFile():
    global m_savefile, m_simuParams;
    plot.saveDataTxt(m_savefile, headers=m_simuParams);
    


O.periodic=True
O.cell.hSize=Matrix3(mul*table.a, 0, 0,
                        0 ,mul*table.a, 0,
                        0, 0, mul*table.a);

def stockesFlow(b):
    global mul;
    
    y0 = mul/2.*table.a;
    dy = (b.state.pos[1]%(mul*table.a)-y0);
    
    
    v0 = Vector3(dy*table.shearRate,0.,0.);
    dv = b.state.vel-v0;
    
    #print("%f, %f, %f, %f"%(y0,dy,v0,dv));
    
    O.forces.setPermF(b.id, -6.*math.pi*table.eta*table.a*dv);
    O.forces.setPermT(b.id, -8.*math.pi*table.eta*table.a**3*(b.state.angVel*(2.*math.pi)-Vector3(0.,0.,-0.5*table.shearRate)));
    
def applyStockesFlow():
    stockesFlow(O.bodies[0]);
    stockesFlow(O.bodies[1]);

# ElectrostaticMat
mat = O.materials.append(CohFrictMat(density=table.rho,young=table.young,poisson=0.3,frictionAngle=radians(60)))

# add two spheres
O.bodies.append([
        sphere(center=(mul/2.*table.a,(mul/2)*table.a,mul/2.*table.a),radius=table.a,material=mat),
        sphere(center=((mul/2.-6)*table.a,(mul/2.+dist)*table.a,mul/2.*table.a),radius=table.a,material=mat)]);

law = Law2_ScGeom_ImplicitLubricationPhys(  activateTangencialLubrication=True,
                                            activateTwistLubrication=True,
                                            activateRollLubrication=True,
                                            theta=1.,
                                            resolution=2,
                                            MaxDist=4.0,
                                            SolutionTol=1.e-9,
                                            MaxIter=100);

O.engines=[ForceResetter(),
           InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=10)]),
           InteractionLoop([Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=10)],
                           [Ip2_FrictMat_FrictMat_LubricationPhys(eta=table.eta,eps=table.roughness)],
                           #[Ip2_ElastMat_ElastMat_LubricationPhys(eta=100,eps=0.)],
                           [law]),
           NewtonIntegrator(damping=0., gravity=(0,0,0),label="newton" ),
           GlobalStiffnessTimeStepper(active=1,timeStepUpdateInterval=100,timestepSafetyCoefficient=0.5, defaultDt=1.e-5, maxDt=1.e-5, label="TimeStepper",viscEl=False),
           
           PyRunner(command='upGraph()',virtPeriod=0.01),
           PyRunner(command='saveToFile()',realPeriod=600),
           PyRunner(command='applyStockesFlow()',iterPeriod=1)
           ]
           
#O.bodies[0].state.blockedDOFs = 'xyzXYZ';
#O.bodies[1].state.blockedDOFs = 'XYZ';

plot.plots = {'t':('Fn','Fnc','Fnl'),'t2':('u','ue'), 'pos1_x':('pos1_y')};
plot.plot(subPlots=True);
O.bodies[0].state.angVel = (0,0.,-0.5*table.shearRate/(2*math.pi));
O.bodies[1].state.angVel = (0,0.,-0.5*table.shearRate/(2*math.pi));
O.saveTmp();

if 'YADE_BATCH' in os.environ :
    for i in range(6,-20,-1):
        d = 10.**(i/10);
        for theta in range(5,91,5):
            O.bodies[1].state.vel = (d*sin(theta)*table.shearRate*table.a,0.,0.);
            O.bodies[1].state.pos = Vector3((mul/2.-6)*table.a,(mul/2.+d*sin(theta))*table.a,(mul/2+d*cos(theta))*table.a);
            O.run(0, True);
            O.loadTmp();
