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

readParamsFromTable(roughness=0, eta=10, a=1.e-2, rho=1e3, g=10., young=5.e7, u0=1.)

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_sediment";

m_tau = table.eta/(table.a*table.rho*table.g);
m_tprim = m_tau + math.sqrt(m_tau**2 + 2.*table.u0*table.a/(table.g));

m_tmax = max(200.*m_tau, 10.*m_tprim);

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)
        
        if O.time > m_tmax:
            SaveAndQuit();
		
def saveToFile():
    global m_savefile, m_simuParams;
    plot.saveDataTxt(m_savefile, headers=m_simuParams);

# 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=(0,0,0),radius=table.a,material=mat,fixed=True),
        sphere(center=(0.0,(table.u0+2)*table.a,0),radius=table.a,material=mat)]);

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

O.engines=[ForceResetter(),
           InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=3)]),
           InteractionLoop([Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=3)],
                           [Ip2_FrictMat_FrictMat_LubricationPhys(eta=table.eta,eps=table.roughness)],
                           #[Ip2_ElastMat_ElastMat_LubricationPhys(eta=100,eps=0.)],
                           [law]),
           NewtonIntegrator(damping=0., gravity=(0,-table.g,0),label="newton" ),
           GlobalStiffnessTimeStepper(active=0,timeStepUpdateInterval=100,timestepSafetyCoefficient=0.8, defaultDt=1e-10,label="TimeStepper",viscEl=False),
           
           PyRunner(command='upGraph()',virtPeriod=m_tau/10),
           PyRunner(command='saveToFile()',realPeriod=600),
           #PyRunner(command='trajectoire()',iterPeriod=1)
           ]

O.dt = 1.e-7;
#O.bodies[1].state.vel = (0,-1e-2,0);
#O.bodies[1].state.blockedDOFs = 'xyzXYZ';

plot.plots = {'t':('Fn','Fnc','Fnl'),'t3':('Fsc_x','Fsc_y', 'Fsl_x', 'Fsl_y'), 't2':('u','ue'), 't4':('nc_yy','nl_yy')};
plot.plot(subPlots=True);
O.saveTmp();
        
#O.run(1);
#O.run(long(0.302/O.dt));

if 'YADE_BATCH' in os.environ:
    O.run();
    waitIfBatch();
