#!/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", "Box"])
    x = np.arange(rout, -rout, -2*r*2**0.5)
    x, y = np.meshgrid(x, x)
    R2 = x**2 + y**2
    keep = R2 < (rout - r)**2
    x = x[keep]
    y = y[keep]
    for i in range(x.shape[0]) :
        p.add_particle((x[i], y[i]), r, r**2 * np.pi * rhop);
    
    p.position()[:, 1] += 0.45
    p.write_vtk(filename,0,0)


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

t = 0
ii = 0

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

#physical parameters
g =  -9.81
rho = 1200
rhop = 2400
nu = 0.8e-4
V = 0.5 # todo : estimate V base on limit velocity
print('V',V)
tEnd = 100

#numerical parameters
lcmin = 0.0001 # approx r*100 but should match the mesh size
dt = 5e-4*10
alpha = 2.5e-2
epsilon = alpha*lcmin**2 /nu
print('epsilon',epsilon)

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

genInitialPosition(outputdir, r, 1e-3, rhop)

p = scontact2.ParticleProblem()
p.read_vtk(outputdir,0)

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

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

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

ii = 0
tic = time.clock()
forces = g*p.mass()
fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
while t < tEnd : 
    if (ii%5==0 and ii != 0):
       fluid.adapt_mesh(1e-3,1e-4,10000)
       fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
    else :
       fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
    fluid.implicit_euler(dt)
    forces = fluid.compute_node_force(dt)
    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)  
    t += dt
    if ii %outf == 0 :
        ioutput = int(ii/outf) + 1
        p.write_vtk(outputdir, ioutput, t)
        #p.write(outputdir+"/part-%05d" % ioutput)
        fluid.export_vtk(outputdir, t, ioutput)
    ii += 1
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.clock() - tic))