#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 14 14:49:05 2019

@author: margauxboxho
"""

# MigFlow - Copyright (C) <2010-2018>
# <Universite catholique de Louvain (UCL), Belgium
#  Universite de Montpellier, France>
# 	
# List of the contributors to the development of MigFlow: see AUTHORS file.
# Description and complete License: see LICENSE file.
# 	
# This program (MigFlow) is free software: 
# you can redistribute it and/or modify it under the terms of the GNU Lesser General 
# Public License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
# 
# You should have received a copy of the GNU Lesser General Public License
# along with this program (see COPYING and COPYING.LESSER files).  If not, 
# see <http://www.gnu.org/licenses/>.

#!/usr/bin/env python

# TESTCASE DESCRIPTION
# Bidimensional particles sedimentation in fluid


from migflow import fluid
from migflow import scontact

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


def genInitialPosition(filename, r, H, ly, lx, rhop, compacity) :
    """Set all the particles centre positions and create the particles objects to add in the computing structure
    
    Keyword arguments:    
    filename -- name of the output file
    r -- max radius of the particles
    H -- domain height
    ly - particles area height
    lx -- particles area width
    rhop -- particles density        
    """
    # Particles structure builder
    p = scontact.ParticleProblem(2)
    # Load mesh.msh file specifying physical boundaries names
    p.load_msh_boundaries("mesh.msh", ["Top", "Lateral","Bottom"])

    #Definition of the points where the particles are located
    #x = np.arange(-lx/2+r, lx/2-r, 2*r)
    #y = np.arange(0-r, -ly+r, -2*r)
    #x, y = np.meshgrid(x, y)
    #x = x.flat
    #y = y.flat
    
    N = 4*compacity*(lx*ly)/(np.pi*r**2)
    e = 2*ly/(N)**(1./2.)

    # Definition of the points where the particles are located
    shift = 0.1; alpha = 30*np.pi/180; 
    lx = lx*np.cos(alpha)
    
    x = np.arange(lx/2+25*r, -lx/2-25*r, -e)
    y = np.arange(-3*r-shift, -ly+3*r-shift, -e)
    
    #shift = 0.26;
    #x = np.arange(lx/2+50*r, -lx/2-50*r, -e)
    #y = np.arange(2*H-3*r-shift, 2*H-ly+3*r-shift, -e)
    x, y = np.meshgrid(x, y)
    
    
    x = x-y*np.tan(alpha); 
    y = y; 
    
    x = x.flat
    y = y.flat
    
    
    # Add a grain at each centre position
    for i in range(len(x)) :
        #p.add_particle((x[i], y[i]), r, r**2 * np.pi * rhop)
        p.add_particle((x[i]+2*random.uniform(-e/2+r,e/2-r), y[i]+2*random.uniform(-e/2+r,e/2-r)), r, r**2 * np.pi * rhop)
    p.write_vtk(filename,0,0)
    

filename = "/Volumes/Margaux/Memoire_Code/testcases/Boycott/"
outputdir = filename + "output"
if not os.path.isdir(outputdir) :
    os.makedirs(outputdir)

# Physical parameters
g = -9.81                                       # gravity
r = 1e-3                                        # particles radius os sand seed
rhop = 1600                                     # particles density as sand
rho = 1000                                      # fluid density
nu = 1e-6                                       # kinematic viscosity
mu = rho*nu;        
compacity = 0.2

# Numerical parameters
outf = 5                                       # number of iterations between output files
dt = 0.5e-3                                    # time step
tEnd = 3                                       # final time

# Geometrical parameters
ly = 0.8                                        # particles area height
lx = 0.1                                        # particles area widht
H = 0.6                                         # domain height

#
# PARTICLE PROBLEM
#
# Initialise particles
#genInitialPosition(outputdir, r, H, ly, lx, rhop, compacity)
p = scontact.ParticleProblem(2)
p.read_vtk("/Volumes/Margaux/Memoire_Code/testcases/Boycott/output_no_slip_wall_Ref",0)

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

U_0 = 2*9.81*(rhop-rho)*r**2/(9*mu); 
phi = compacity
f_phi = ((1-phi)**2)/((1+phi**(1/3))*np.exp(5/3*phi/(1-phi))); 

U_0 = 0.8633; 
P_0 = 255;
Re = rho*U_0*r/mu; 
N_0 = len(p.position());
e_target = 1e-7
print(" ======= The initial number of grains is %.5f [-]=======" %N_0)
print(" ======= The initial velocity is %.5f [m/s]=======" %U_0)
print(" ======= The initial pressure is %.5f [Pa]=======" %P_0)
print(" ======= The Reynolds number is %.3f [Pa]=======" %Re)

# Initial time and iteration
t         =  0
ii        =  0
jj        =  0
delta     = 10

#
# FLUID PROBLEM
#
fluid = fluid.FluidProblem(2,g,[nu*rho],[rho], petsc_solver_type="-pc_type lu -ksp_max_it 50")
# Set the mesh geometry for the fluid computation
fluid.load_msh("mesh.msh")
fluid.set_wall_boundary("Bottom", velocity=[0,0])
fluid.set_wall_boundary("Lateral", velocity=[0,0])
fluid.set_wall_boundary("Top",pressure=0, velocity=[0,0])
# Set location of the particles in the mesh and compute the porosity in each computation cell
fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity(), p.contact_forces())
fluid.export_vtk(outputdir,0,0)

tic = time.time()
fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity(), p.contact_forces())
  
#
# COMPUTATION LOOP
#
while t < tEnd :
    if (ii==jj*delta):
        print("Save state at %i" %(ii))
        fluid.export_vtk(outputdir,0,0)
        p.write_vtk(outputdir,0,0)
        
    # Fluid solver
    fluid.implicit_euler(dt)
    
    if (ii==0):
        fluid.adapt_mesh(60*r, 5*r, 6602*2, e_target, 1, U_0, P_0)
        
    if (ii%(delta*(jj+1))!=0): 
        fluid.adapt_mesh(60*r, 5*r, 6602*2, e_target, 0, U_0, P_0) 
    
    # Adaptation of the mesh and loading the past save state
    if (ii%(delta*(jj+1))==0 and ii != 0):  
        print("Adapt mesh at %i" %(ii))
        # Adapt the mesh by creating the lc.pos file
        fluid.adapt_gen_mesh(60*r, 5*r, 6602*2, e_target, U_0, P_0)
        ii = ii-delta
        t = t-dt*delta
        jj += 1
        print("ii=%i, t=%i, jj=%i" %(ii,t,jj))
        # load the save past state 
        folder = outputdir+"/fluid_00000.vtu"
        fluid.import_vtk(folder)
        p.read_vtk(outputdir,0)
        fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity(), p.contact_forces())
        # create the .msh on the adapt/lc.pos
        fluid.adapt_load_mesh("adapt/mesh.msh")
        #fluid.implicit_euler(dt)
        print("Restart computation %i iterations back at %i" %(delta, ii)) 
        # Initialize the min_h_target attribut of the Mesh structure
        fluid.adapt_mesh(60*r, 5*r, 6602*2, e_target, 1, U_0, P_0) 
        
    # Computation of the new velocities
    forces = fluid.compute_node_force(dt)
    vn = p.velocity() + forces*dt/p.mass()
    vmax = np.max(np.hypot(vn[:, 0], vn[:, 1]))
    # Number of contact sub time step
    nsub = max(1, int(np.ceil((vmax*dt*4)/min(p.r()))))
    print("NSUB", nsub,"VMAX",vmax, "VMAX * dt", vmax * dt, "r", min(p.r()))
    # NLGS iterations
    for i in range(nsub) :
        tol = 1e-8                                #numerical tolerance for particles intersection
        p.iterate(dt/nsub, forces, tol)

    # Localisation of the particles in the fluid
    fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity(), p.contact_forces())
    t += dt

    # Output files writting
    if ii %outf == 0 :
        ioutput = int(ii/outf) + 1
        p.write_vtk(outputdir, ioutput, t)
        fluid.export_vtk(outputdir, t, ioutput)
    ii += 1
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.time() - tic))

"""

while t < tEnd :
    # Save the state of the fluid at certain iteration 
    if (ii==jj*delta):
        print("Save state at %i" %(ii))
        fluid.export_vtk(outputdir,0,0)
        p.write_vtk(outputdir,0,0)
        
    # Fluid solver
    fluid.implicit_euler(dt)
    
    if (ii==0):
        fluid.adapt_mesh(500*r, 8*r, 6602*2, 0, 1)
        
    if (ii%(delta*(jj+1))!=0): 
        fluid.adapt_mesh(500*r, 8*r, 6602*2, 0, 0)
    
    # Adaptation of the mesh and loading the past save state
    if (ii%(delta*(jj+1))==0 and ii != 0):  
        print("Adapt mesh at %i" %(ii))
        #fluid.adapt_gen_mesh(500*r, 8*r, 6602*2, 1e-6, "adapt/lc.pos")
        # Adapt the mesh by creating the lc.pos file
        fluid.adapt_mesh(500*r, 8*r, 6602*2, 1, 0)
        ii = ii-delta
        t = t-dt*delta
        jj += 1
        print("ii=%i, t=%i, jj=%i" %(ii,t,jj))
        # load the save past state 
        fluid.import_vtk("output/fluid_00000.vtu")
        p.read_vtk(outputdir,0)
        fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
        # create the .msh on the adapt/lc.pos
        fluid.adapt_load_mesh("adapt/mesh.msh")
        #fluid.implicit_euler(dt)
        print("Restart computation %i iterations back at %i" %(delta, ii)) 
        # Initialize the min_h_target attribut of the Mesh structure
        fluid.adapt_mesh(500*r, 8*r, 6602*2, 0, 1)
    
    # Computation of the new velocities
    forces = fluid.compute_node_force(dt)
    # print(forces)
    vn = p.velocity() + forces*dt/p.mass()
    vmax = np.max(np.hypot(vn[:, 0], vn[:, 1]))
    # Number of contact sub time step
    nsub = max(1, int(np.ceil((vmax*dt*4)/min(p.r()))))
    print("NSUB", nsub,"VMAX",vmax, "VMAX * dt", vmax * dt, "r", min(p.r()))
    # NLGS iterations
    for i in range(nsub) :
        tol = 1e-8                                #numerical tolerance for particles intersection
        p.iterate(dt/nsub, forces, tol)

    # Localisation of the particles in the fluid
    fluid.set_particles(p.mass(), p.volume(), p.position(), p.velocity())
    t += dt

    # Output files writting
    if ii %outf == 0 :
        ioutput = int(ii/outf) + 1
        p.write_vtk(outputdir, ioutput, t)
        fluid.export_vtk(outputdir, t, ioutput)
    ii += 1
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.time() - tic))

"""