# %%
# Lock-exchange testcase
# ======================
# This example illustrates the classical lock-exchange problem in a closed
# basin. Two fluids of different density are initially separated in a domain;
# when the barrier is removed, the dense fluid flows along the bottom and the
# light fluid flows above in the opposite direction, forming gravity currents
# and a propagating density front. It is commonly used to test an ocean model’s
# ability to represent stratified flow, sharp interfaces, and mixing.

import os
import time

import numpy as np
from slim4 import mpi, slim, slim2d, slim3d, pre3d
from slim4.tools import visualize_prepro_data

# %%
# We optionally write preprocessing checks to disk. These files can be useful
# when validating the bathymetry, vertical extrusion, and initial tracer fields.

check_preprocessing = False


# %%
# First, we load the 2D mesh and define the bathymetry and free surface used to
# build the 3D domain.

print("Loading hydro setup ...")
mesh2d = slim2d.Mesh("lock_exchange_structured.msh")
num_2d_elements = mesh2d.n_triangles

prepro_dir = "prepro"
os.makedirs(prepro_dir, exist_ok=True)
print(f"Prepro dir: {prepro_dir}")

bathymetry_2d = slim.DataStatic.create_constant(mesh2d, [20])
free_surface_2d = slim.DataStatic.create_constant(mesh2d, [0])
slim.DataFile.write(
    f"./{prepro_dir}/bathymetry", [("Depth", bathymetry_2d.eval_mesh())], "DG"
)


# %%
# Next, we define the vertical discretization. This example uses sigma layers
# We extrude the 2D mesh vertically to obtain the 3D computational mesh.
n_layers = 20
mesh, z0, bathymetry_2d = pre3d.extrude(mesh2d, bathymetry_2d, pre3d.VMapProportional(np.ones(n_layers)))


# %%
# Now we create the 3D hydrostatic problem and add the tracers used in the
# testcase. Temperature controls the density anomaly, while salinity is kept as
# a passive background field in this configuration.

g = 9.80616
rho_0 = 1000
num_3d_elements = mesh.z_layers[-1]

problem = slim3d.Problem(mesh, g, bathymetry_2d, rho_0)
temperature_equation = problem.add_tracer(1, "temperature")
salinity_equation = problem.add_tracer(1, "salinity")
problem.set_wall_boundary("Wall", slip=False)


# %%
# We initialize the prognostic fields and moving vertical coordinates used by
# the 3D model.

z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = problem.problemUV.new_vector()
temperature = temperature_equation.new_vector()
salinity = salinity_equation.new_vector()


# %%
# The initial temperature field defines the lock exchange: colder, denser water
# is placed on the left side of the domain, while warmer, lighter water fills
# the rest of the basin.

temp_init = np.ones((1, 6, num_3d_elements)) * 30
x_threshold = 32e3
left_columns = np.all(mesh2d.x[0, :, :] <= x_threshold, axis=0)
left_triangles = np.arange(num_2d_elements)[left_columns]

for triangle in left_triangles:
    column = np.arange(mesh.z_layers[triangle], mesh.z_layers[triangle + 1])
    temp_init[0, :, column] = 5

temperature.set(temp_init)
slim.DataFile.write(f"./{prepro_dir}/temperature", [("Temperature", temp_init)], "DG")


# %%
# Salinity is initialized uniformly. The density is then linked to temperature
# through a linear equation of state.

salinity_initial = np.ones((1, 6, num_3d_elements)) * 35
salinity.set(salinity_initial)

problem.set_state_equation_linear([(temperature, -0.2)])


# %%
# Optionally, we write preprocessing figures and 3D fields to help verify that
# the bathymetry and initial condition were generated correctly.

if check_preprocessing:
    visualize_prepro_data(
        f"./{prepro_dir}/bathymetry",
        "Depth",
        f"./{prepro_dir}/check/bathymetry",
        mesh2d,
    )
    mesh.create_xdmf(z0, f"./{prepro_dir}/check/temperature")
    mesh.append_to_xdmf(
        f"./{prepro_dir}/check/temperature", 0, 0, z0, [("temperature", temp_init)]
    )


# %%
# Next, we define the viscous and diffusive parameters. This testcase uses no
# explicit horizontal tracer diffusion, a fixed vertical viscosity, and a small
# horizontal viscosity.

vertical_viscosity = np.full((1, 1, mesh.n_elements), 1e-4)
problem.problemUV.set_vertical_viscosity(
    slim.ScalarArrayDevice.create_from_host(vertical_viscosity)
)

horizontal_viscosity = slim.DataStatic.create_constant(mesh, value=1e-2)
problem.problemUV.set_horizontal_viscosity(horizontal_viscosity, "direct")


# %%
# We prepare the output files and choose the time-stepping parameters.

nu_h_value = horizontal_viscosity.eval_mesh()[0, 0, 0]
print(
    f" -> Running lock-exchange testcase with nu_h = {nu_h_value:.3g} m^2/s and {n_layers} vertical layers"
)

output_dir = "out/"
path2d = output_dir + "slim2D"
path3d = output_dir + "slim3D"
mesh2d.create_xdmf(path2d)
mesh.create_xdmf(z0, path3d)

problem.problemUV.update_density(z, problem.get_gpu_tmpvars().rho)

t_end = 17 * 3600
export_dt = 900

rho = problem.get_gpu_tmpvars().rho
problem.problemUV.update_density(z, rho)
dt3d, dt2d = problem.compute_cfl(2, rho.get())

nsub = int(np.ceil(export_dt / dt3d))
dt = export_dt / nsub
M = dt / dt2d
nstep = int(np.ceil(t_end / dt))
export_stride = int(np.ceil(export_dt / dt))

print(f"dt3d = {dt:.3f}, dt2d = {dt2d:.2f}, M = {M:.2f}, nstep = {nstep}")


# %%
# Finally, we run the simulation and export the solution at regular intervals.
# The exported fields include density diagnostics in addition to temperature and
# salinity.

t = 0
iout = 0
export_kwargs = dict(turbulence=False, Huv=False, rho=True)
problem.do_export(
    path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
)

start_time = time.time()
loop_start = start_time
iteration = 0
while t + dt <= t_end:
    t = slim3d.iterate_impl(problem, t, dt, M, 1, z, uv, [temperature, salinity])
    # t = slim3d.iterate_midpoint(problem, t, dt, M, 1, z, uv, [temperature, salinity])

    export_index, remainder = divmod(iteration + 1, export_stride)
    if remainder == 0:
        iout += 1
        problem.do_export(
            path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
        )
        uv_max = mpi.max_scalar(np.linalg.norm(uv.get(), axis=0).max())
        w_max = mpi.max_scalar(np.abs(problem.get_gpu_tmpvars().w_gpu.get()).max())
        print(
            f"Output {iout} at t = {t / 3600:.2f}h, "
            f"uv_max = {uv_max:.2e}, w_max = {w_max:.2e}, "
            f"{slim.si_format((time.time() - loop_start) / (iteration + 1), 3)}s/iter"
        )

    iteration += 1

print(f"Hydro finished in {(time.time() - start_time) / 60:.2f} minutes")
# %%