#!/usr/bin/env python
from marblesbag import fluid as fluid
from marblesbag import scontact2

import numpy as np
import os
import time
import shutil
import random

def genInitialPosition(filename, r, rout, rhop) :
    p = scontact2.ParticleProblem()
    p.load_msh_boundaries("mesh.msh", ["Top", "Bottom","Left","Right"])
        
    p.add_particle((0, 0), r, r**2 * np.pi * rhop);

    p.position()[:, 1] += 0.45
    p.write(filename)



outputdir = "output"
if not os.path.isdir(outputdir) :
    os.makedirs(outputdir)

t = 0
ii = 0

r=0.36e-3/2
p = scontact2.ParticleProblem()
#R = np.random.uniform(45e-06, 90e-06, len(x))

#physical parameters
g =  -9.81
rho = 950
rhop = 2440
nu = 174e-6#/(10**0.5)
V = 0.5 # todo : estimate V base on limit velocity
print('V',V)
tEnd = 10000

#numerical parameters
lcmin = 0.005 # approx r*100 but should match the mesh size
dt = 1e-2
alpha = 2.5e-4
epsilon = alpha*lcmin**2 /nu#2e-2*c*h*2#2e-2*c*h # ?? not sure ??1e-5
autoEpsilon = False
print('epsilon',epsilon)

shutil.copy("mesh.msh", outputdir +"/mesh.msh")
#scontact2Interface.MeshLoader(p, "funnel.msh", ("Funnel", "Top", "Bottom", "Lateral"))
p.write(outputdir+"/part-00000")

genInitialPosition(outputdir + "/part-00000", r, 1e-3, rhop)

p = scontact2.ParticleProblem(outputdir+"/part-00000")

print("r = %g, m = %g\n" %  (p.r()[0], p.mass()[0]))
print("RHOP = %g" % rhop)
outf = 1
outf1 = 100000

strong_boundaries = [("Top",2,0.),("Top",1,0.),("Top",0,0.),("Bottom",1,0.),("Bottom",0,0.),("Left",0,0.),("Left",1,0.),("Right",0,0.),("Right",1,0.)]
fluid = fluid.fluid_problem("mesh.msh",g,nu*rho,rho,epsilon,autoEpsilon,strong_boundaries)

fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
fluid.export(outputdir,0,0)

ii = 0
    

tic = time.clock()
forces = g*p.mass()
while t < tEnd : 
   # if (ii%50==0 and ii != 0):
    #   fluid.adapt_mesh(0.01,50,4e-5,autoEpsilon)
    forces = fluid.compute_node_force(dt,r*100)
    vn = p.velocity() + forces * dt / p.mass()
    vmax = np.max(np.hypot(vn[:, 0], vn[:, 1]))
    nsub = max(1, int(np.ceil((vmax * dt * 4)/min(10*p.r()))))
    print("NSUB", nsub,"VMAX",vmax, "VMAX * dt", vmax * dt, "r", min(p.r()))
    for i in range(nsub) :
        p.iterate(dt/nsub, forces)  
    fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
    fluid.implicit_euler(dt)
    t += dt
    if ii %outf == 0 :
        ioutput = int(ii/outf) + 1
        p.write_vtk(outputdir, ioutput, t)
        p.write(outputdir+"/part-%05d" % ioutput)
        p.write_boundary_vtk(outputdir, ioutput, t)
        fluid.export(outputdir, t, ioutput)
    ii += 1
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.clock() - tic))