# %%
# 2D Falling Disc in a Viscous Fluid
# ==================================
# This example simulates the settling of a single solid disc in a viscous
# fluid under gravity in two dimensions. The fluid mesh is refined around
# the moving particle to accurately capture the hydrodynamic interactions.
# %%
# Keywords
# --------
# FEM, DEM, Sedimentation, Mesh Adaptation
from migflow import fluid, gmsh_io, volume_coupling as vc
from migflow import scontact

import numpy as np
import os
import sys
import time
import gmsh
import shutil

gmsh.initialize()

# %%
# Output Directory
# ----------------
outputdir = "output_2d_falling_disc" if len(sys.argv) < 2 else sys.argv[1]
clscale = float(sys.argv[2]) if len(sys.argv) > 2 else 20 #reduce the meshsize for convergence analysis

# %%
# Parameters
# ----------

L = 0.04
H = 0.4
r = 5e-3 / 2
rho = 996
nu = 8e-7
mu = rho * nu
rhop = 1.01 * rho
xp = np.array((0, H / 2 - 4 * r))
lc = L * 1e-3 * clscale
lc_max = lc * 10
lc_min = lc
Re = 156
vRef = Re * mu / (rho * 2 * r)
vRef = 0.02501
print("vRef = %g" % vRef)


# %%
# Mesh generation
# ---------------
def size_callback(x, y):
    d = np.abs(x - xp[0])
    d = d / L

    distmin = 75 * r
    alpha = np.clip((d - distmin) / (10 * distmin), 0, 1)
    size = lc_min * (1 - alpha) + lc_max * alpha
    return size
    # linear
    return np.maximum(lc_max * np.maximum(d - 0.001, 0.0), lc_min)
    if d < 0.1:
        return lc_min
    return lc_max


def gen_mesh(h, w, mesh_size, origin=np.array([0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")

    p1 = gmsh.model.geo.add_point(-w / 2, +h / 2, 0, mesh_size)
    p2 = gmsh.model.geo.add_point(-w / 2, -h / 2, 0, mesh_size)
    p3 = gmsh.model.geo.add_point(+w / 2, -h / 2, 0, mesh_size)
    p4 = gmsh.model.geo.add_point(+w / 2, +h / 2, 0, mesh_size)
    l1 = gmsh.model.geo.add_line(p1, p2)
    l2 = gmsh.model.geo.add_line(p2, p3)
    l3 = gmsh.model.geo.add_line(p3, p4)
    l4 = gmsh.model.geo.add_line(p4, p1)
    line_loop = gmsh.model.geo.add_curve_loop([l1, l2, l3, l4])
    surface = gmsh.model.geo.add_plane_surface([line_loop])
    gmsh.model.geo.synchronize()
    gmsh.model.add_physical_group(1, [2], name="Bottom")
    gmsh.model.add_physical_group(1, [4], name="Top")
    gmsh.model.add_physical_group(1, [1, 3], name="Lateral")
    gmsh.model.add_physical_group(2, [1], name="domain")
    gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: size_callback(x, y))
    gmsh.option.set_number("Mesh.Algorithm", 1)
    gmsh.model.mesh.generate(2)


gen_mesh(H, L, lc)

t = 0
ii = 0

print("Start with r=%g" % (r / 2))
shutil.rmtree(outputdir, True)
if not os.path.isdir(outputdir):
    os.makedirs(outputdir)
# physical parameters
g = np.array([0, -9.81])  # gravity
tEnd_adim = 10
tEnd = tEnd_adim * 2 * r / vRef  # final time

# %%
# Numerical parameters
# --------------------
cfl = 0.001 / 10
dt = cfl / vRef  # time step
outf = 10

# %%
# Particle problem definition
# ---------------------------
p = scontact.ParticleProblem(2)
p.add_particle(xp, r, r**2 * np.pi * rhop)
p.write_mig(outputdir, 0)

# %%
# Fluid problem definition
# ------------------------
f = fluid.FluidProblem2(g, nu * rho, rho)
gmsh_io.load_msh(f, None)
f.set_symmetry_boundary("Bottom")
f.set_symmetry_boundary("Lateral")
f.set_symmetry_boundary("Top")
f.set_mean_pressure(0)
# Tier A (the disc's geometry on the mesh) and tier B (the legacy mixture
# closure) only: this primes the coupling so the frame written below carries a
# porosity field. The loop rebuilds both, with the Patankar datum, before every
# solve.
geo = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                               omega=p.omega(),
                               contact_forces=p.contact_forces())
vc.set_coupling_geometry(f, geo)
vc.set_coupling_closure(f, clo)

def get_fields(fluid):
    """Derived output fields: the vorticity is what shows the wake of a
    falling disc, so the case writes it beside the primitives."""
    p1_element = fluid.get_p1_element()
    grad_v = fluid.fields_gradient()[:, :2, :]
    # porosity().get() is 1-D: reshape it too, or (n, 1) * (n,) broadcasts to
    # the (n, n) outer product (105 GB on disk once, see 2d_square).
    vorticity = ((grad_v[:, 1, 0] - grad_v[:, 0, 1]).reshape(-1, 1)
                 * fluid.porosity().get().reshape(-1, 1))
    return {
        "pressure": (fluid.pressure(), p1_element),
        "velocity": (fluid.velocity(), p1_element),
        "porosity": (fluid.porosity().get(), p1_element),
        "vorticity": (vorticity, p1_element),
    }


f.write_mig(outputdir, 0, get_fields(f))

# %%
# Computation Loop
# ----------------
t = 0
ii = 0
tic = time.perf_counter()
mass = p.r() ** 2 * np.pi * rhop
dp = []
drag = []
total_x = []
total_y = []
pvelocity_x = []
pvelocity_y = []
while t < tEnd:
    t += dt
    t_adim = t * vRef / (2 * r)
    print("t = %.2g" % t)
    print("t_adim = %.2g" % t_adim)
    # The bodies sit at their end-of-step positions; the disc is free and the
    # implicit_euler below is the Patankar step that advances it, so this
    # generation carries a datum.
    geo = vc.get_particles(f, p.position() + p.velocity() * dt, p.r(),
                           density=vc.body_density(p))
    clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                                   omega=p.omega(),
                                   contact_forces=p.contact_forces())
    dat = vc.get_patankar(geo, clo, dt, g, p.velocity())
    vc.set_coupling_geometry(f, geo)
    vc.set_coupling_closure(f, clo)
    vc.set_coupling_datum(f, dat)
    f.implicit_euler(dt)
    pforce = f.compute_node_force()  # -drag-dp
    if ii % outf == 0:
        f.write_mig(outputdir, t_adim, get_fields(f))
    p.iterate(dt, pforce + g * mass)
    # Output files writting
    if ii % outf == 0:
        p.write_mig(outputdir, t_adim)
    ii += 1
    dp_i = np.array([0.0])
    dp.append(dp_i)
    drag.append(-(pforce[:, 1] + dp_i))
    total_y.append(pforce[:, 1])
    total_x.append(pforce[:, 0])
    pvelocity_y.append(p.velocity()[-1, 1])
    pvelocity_x.append(p.velocity()[-1, 0])
last_velocity = p.velocity()[-1, 1]
print(last_velocity)

dp = np.array(dp).reshape(-1)
drag = np.array(drag).reshape(-1)
weight = g[1] * np.pi * r**2 * rhop * np.ones(dp.shape[0])
total_y = np.array(total_y).reshape(-1)
total_x = np.array(total_x).reshape(-1)

i = np.arange(drag.shape[0])
fmax = np.max(np.abs(total_y))
drag /= fmax
dp /= fmax
total_y /= fmax
total_x /= fmax
weight /= fmax
pvelocity_y = np.abs(np.array(pvelocity_y).reshape(-1)) / vRef
pvelocity_x = np.array(pvelocity_x).reshape(-1) / vRef


simu_time = np.arange(dp.shape[0]) * dt
np.savetxt("time.txt", simu_time)
np.savetxt("drag.txt", drag)
np.savetxt("dp.txt", dp)
np.savetxt("weight.txt", weight)
np.savetxt("total_y.txt", total_y)
np.savetxt("total_x.txt", total_x)
np.savetxt("pvelocity_y.txt", pvelocity_y)
np.savetxt("pvelocity_x.txt", pvelocity_x)

PLOTTING = False
if PLOTTING:
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 2)
    buoyancy = (rhop - rho) * g[1] * np.ones(dp.shape[0]) * np.pi * r**2 / fmax
    ax[0].plot(i, -dp, "red", label="dp")
    ax[0].plot(i, -drag, "blue", label="drag")
    ax[0].plot(i, weight, "green", label="weight")
    ax[0].plot(i, total_y, "--k", label="total")
    ax[0].plot(i, total_x, "--", color="red", label="total")
    ax[0].set_ylim([-2, 2])
    ax[0].legend()

    ax[1].plot(i, pvelocity_y)
    ax[1].plot(i, pvelocity_x)
    ax[1].set_ylim([-0.2, 1.2])
    plt.show()

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_falling_disc --actors fluid particles --fluid-field vorticity --bounds -0.02 0.02 0.1 0.2 --show-edges 1
# %%
# Artifacts
# ---------
# - output_2d_falling_disc/animation.mp4