# MigFlow - Copyright (C) <2010-2022>
# <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
import os
import time
import gmsh
import shutil
import numpy as np
from migflow import fluid
from migflow import scontact
from migflow import time_integration
from pylmgc90 import pre

np.random.seed(0)


def gen_mesh(width, height, mesh_size, origin=np.array([0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")
    gmsh.model.occ.add_rectangle(origin[0], origin[1], 0, width, height)
    gmsh.model.occ.synchronize()

    def get_line(x0, x1, eps=1e-6):
        r = gmsh.model.get_entities_in_bounding_box(
            x0[0] - eps, x0[1] - eps, -eps, x1[0] + eps, x1[1] + eps, eps, 1
        )
        return list(tag for dim, tag in r)

    h, w = height, width
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, 0], origin + [w, 0]), name="Bottom"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, h], origin + [w, h]), name="Top"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, 0], origin + [0, h]), name="Left"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [w, 0], origin + [w, h]), name="Right"
    )
    gmsh.model.add_physical_group(2, [1], name="domain")
    gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: mesh_size)
    # gmsh.option.setNumber("Mesh.Algorithm", 2)  # Frontal-Delaunay for 2D meshes
    gmsh.model.mesh.generate(2)


# Define output directory
outputdir = "output-compact/rayleigh_with_friction"
shutil.rmtree(outputdir, ignore_errors=True)
if not os.path.isdir(outputdir):
    os.makedirs(outputdir)

# Physical parameters
g = np.array([0, -9.81])  # gravity
r = 70e-6  # particles radius
rhop = 1050  # particles density : polystyrene [kg/m^3]
rho = 1.204  # fluid density : air [kg/m^3]
nu = 1.516e-5  # kinematic viscosity
mu = rho * nu
# Numerical parameters
outf = 50  # number of iterations between output files
dt = 2e-5  # time step
tEnd = 1  # final time
# Geometrical parameters
height = 68e-3  # domain height [m]
width = 56e-3  # domain width [m]
depth = 1e-3  # domain depth [m]
mesh_size = height / 50  # mesh size
h = 49.5e-3  # particles area height
w = width  # particles area widht
origin = np.array([-width / 2, -height / 2])
eps = r * 1e-3
# Generate the mesh
gen_mesh(width, height, mesh_size, origin)

# Particle Problem
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(True)
p.load_msh_boundaries(None, ["Top", "Left", "Right", "Bottom"])
p.set_friction_coefficient(0.3)

n_particles = 80000
radii = np.random.uniform(0.9, 1.1, n_particles) * r
[nb_deposited, x] = pre.depositInBox2D(radii, width, height)
x = x.reshape(-1, 2)
x[:, 1] *= -1
offset = np.array([origin[0], -origin[1]])
x = x + offset[None, :]
for xi, ri in zip(x, radii[:nb_deposited]):
    p.add_particle(xi, ri, ri**2 * np.pi * rhop)

# Fluid Problem
f = fluid.FluidProblem(2, g, mu, rho, usolid=True, volume_drag=12 * mu / (depth**2))
f.load_msh(None)
f.set_wall_boundary("Bottom")
f.set_wall_boundary("Left")  # periodicity can be imposed
f.set_wall_boundary("Right")  # periodicity can be imposed
f.set_wall_boundary("Top")
f.set_mean_pressure(0)
omega = np.zeros_like(p.omega())
f.set_particles(
    p.delassus(), p.volume(), p.position(), p.velocity(), omega, p.contact_forces()
)

# Initial time and iteration
t = 0
ii = 0
p.write_mig(outputdir, 0)
f.write_mig(outputdir, 0)
# COMPUTATION LOOP
mass = np.pi * p.r() ** 2 * rhop
tic = time.process_time()
print("Atwood number : ", (rhop - rho) / (rho + rhop))
while t < tEnd:
    print("%i : %.2g/%.2g (cpu %.6g)" % (ii, t, tEnd, time.process_time() - tic))
    f.set_particles(
        p.delassus(), p.volume(), p.position(), p.velocity(), omega, p.contact_forces()
    )
    tic = time.process_time()
    f.implicit_euler(dt, check_residual_norm=1)
    print("implicit : ", time.process_time() - tic)
    fext = f.compute_node_force(dt) + g * mass
    tic = time.process_time()
    time_integration._advance_particles(p, fext, dt, 1, 1e-6 * r)
    print("NSCD : ", time.process_time() - tic)
    t += dt
    # Output files writting
    if ii % outf == 0:
        f.write_mig(outputdir, t)
        p.write_mig(outputdir, t)
    ii += 1