# %%
# Dome testcase
# =============
# This example illustrates a classical 3D coastal-ocean testcase where a
# river inflow injects water and a passive tracer into the domain. As the
# inflow develops, the Coriolis force deflects the jet and produces the
# characteristic dome-like circulation pattern.
#
# The setup combines open boundaries, prescribed inflow conditions, and a
# fully 3D hydrostatic model with tracers.

from slim4 import slim, slim2d, slim3d, pre3d, mpi
from slim4.slim import si_format
import numpy as np
import time

tic = time.time()
check3d = True

# %%
# First, we load and partition the 2D mesh, then construct the bathymetry.
# The bathymetry is defined from a simple linear slope and clipped so that the
# water depth remains within a prescribed range.

mesh2d = slim2d.Mesh(slim2d.partition_mesh("dome.msh"))  # run_coarse.msh, run_fine.msh
bb = -0.01 * mesh2d.x[1] + 6600
bb[bb < 600] = 600
bb[bb > 3600] = 3600
bath2d = slim.DataStatic(mesh2d, "DG", 1, bb)
nt = mesh2d.x.shape[-1]

# %%
# Next, we extrude the 2D mesh into a 3D mesh. In this example we use a
# sigma discretization, but a z-coordinate extrusion is also shown as an
# alternative in the commented line below.

nl = 24
z_l = np.linspace(0, -3600, nl, endpoint=True)
vmap = pre3d.VMapProportional(np.ones(nl)) # sigma coordinates
# vmap = pre3d.VMapZ.create_from_vmap(vmap, bath2d) # z coordinates
mesh, z0, bath2d = pre3d.extrude(mesh2d, bath2d, vmap)

n_elem = mesh.z_layers[-1]
X = mesh.create_coordinate_vector(z0).T


# %%
# We now define the boundary data used at the open boundary. The inflow is
# prescribed through a vertically structured velocity profile as described in
# `Legg et al. (2006) <https://doi.org/10.1016/j.ocemod.2004.11.006>`_ and define
# the plume entering the bassin.

eta_open = slim.DataStatic.create_constant(mesh2d, [0])
c_in = slim.DataStatic.create_constant(mesh, [1])
u_open = slim.DataStatic(mesh, type="DG", nfields=2)
t_in = slim.DataStatic(mesh, type="DG", nfields=1)

b0 = 0.019
h0 = 300
u0 = np.sqrt(h0 * b0)
xw = X[..., 0] - 800e3
he = 600
Rim = 1 / 3
N2 = 2.3e-3**2
f = 1e-4
g = 9.80616
rho0 = 1000


# %%
# The helper function below defines the vertical shape of the inflow profile.
# It is then used to construct the velocity, density and temperature signals
# imposed at the inflow boundary.


def F(zs):
    zs = np.array(zs)
    out = np.zeros_like(zs)
    high = np.where(zs >= 1 / 5)[0]
    mid = np.where((zs > -1 / 7) & (zs < 1 / 5))[0]
    low = np.where(zs <= -1 / 7)[0]
    out[high] = 1
    out[mid] = 1 / 3 * zs[mid] / (zs[mid] + 1) + 0.5
    out[low] = 0
    return out


# modified z coordinate
z = X[..., 2]
h_func = h0 * np.exp(-xw * f / u0)
z_s = (z - h_func + he + 1e-3) / (h_func)

# inflow velocity
u_in = -u0 * np.exp(-xw * f / u0) * (1 - F(z_s))

# inflow buoyancy
b_in = np.minimum(-b0 * (1 - F(z_s)), b0 + N2 * z)

# inflow density
rho_in = -rho0 / g * b_in + rho0

# background temperature in bassin
t_stratif = (-rho_in + rho0) / 0.2 + 20

tri_in, _ = mesh2d.boundary_nodes_dg(["In"])
prims_in = np.concatenate(
    [mesh.z_layers[tri] + np.arange(mesh.n_layers[tri]) for tri in tri_in]
)

u_in_help = np.zeros((n_elem, 6))
u_in_help[prims_in, :] = u_in[prims_in, :]
u_open_fields = u_open.eval_mesh()
u_open_fields[1] = u_in_help.T[None, :, :]
u_open.set(u_open_fields)

hu_open = slim3d.DataHMultiply(u_open, eta_open, bath2d, z0)
# hu_open.set_ramp(0, 2 * 3600 * 24)
# eta_open.set_ramp(0, 2 * 3600 * 24)

temp0 = (1 / 360) * X[..., 2] + 20
t_in_help = np.zeros_like(temp0)
t_in_help[prims_in, :] = t_stratif[prims_in, :]
t_in_help[~prims_in, :] = temp0[~prims_in, :]
t_in.set(t_in_help.T[None, :, :])


if check3d:
    mesh.create_xdmf(z0, "./prepro/check/temperature")
    mesh.append_to_xdmf(
        "./prepro/check/temperature", 0, 0, z0, [("temperature", t_in.eval_mesh())]
    )
    mesh.create_xdmf(z0, "./prepro/check/u_in")
    u_in_xdmf = np.zeros((2, 6, n_elem))
    u_in_xdmf[0] = 0.0
    u_in_xdmf[1] = u_in_help.T[None, :, :]
    mesh.append_to_xdmf("./prepro/check/u_in", 0, 0, z0, [("u", u_in_xdmf)])

# %%
# We can now create the 3D problem, add the temperature and passive-tracer
# equations, and prescribe the boundary conditions. The inflow boundary carries
# both temperature and tracer, while the outflow is left open.

problem = slim3d.Problem(mesh, g, bath2d, rho0)
temperature_equation = problem.add_tracer(1, "temperature", limiter_type="past_bound")
tracer_equation = problem.add_tracer(1, "tracer", limiter_type="past_bound")

t_out = slim.DataStatic(mesh, type="DG", nfields=1, value=temp0.reshape(-1, 6).T)
c_out = slim.DataStatic.create_constant(mesh, [0])

problem.set_wall_boundary("Wall")
problem.set_open_boundary(
    "In",
    eta=eta_open,
    u=hu_open,
    t={"temperature": t_in, "tracer": c_in},
    uv_is_transport=True,
)
problem.set_open_boundary("Out")
# problem.set_open_boundary("Out", eta=eta_open, t={"temperature": t_out, "tracer": c_out})
# problem.add_turbulence()


# %%
# Next, we define the physical parameters of the momentum and tracer equations:
# bottom drag, Coriolis force, viscosity, and tracer diffusivity.

gamma = slim.DataStatic.create_constant(mesh2d, [2e-3])
problem.problemUV.set_dissipation(gamma, log_profile=False)

cor = slim.DataStatic.create_constant(mesh2d, 1.0e-4)
problem.problemUV.set_coriolis(cor)

nu_h = slim.DataStatic.create_constant(mesh, 5.0)
problem.problemUV.set_horizontal_viscosity(nu_h)
nu_v = slim.ScalarArrayDevice.create_from_host(5e-2 * np.ones([1, 6, n_elem]))
problem.problemUV.set_vertical_viscosity(nu_v)

kappa_h = slim.DataStatic.create_constant(mesh, 10)
temperature_equation.set_horizontal_diffusivity(kappa_h)
kappa_v = slim.ScalarArrayDevice.create_from_host(0.0 * np.ones([1, 6, n_elem]))
temperature_equation.set_vertical_diffusivity(kappa_v)


# %%
# We initialize the prognostic fields: the passive tracer, the temperature
# tracer, the moving vertical coordinates, and the 3D velocity field.
#
# The passive tracer starts at zero in the domain and is injected through the
# inflow boundary. The temperature field starts from a prescribed background
# profile. A linear equation of state can be enabled if temperature should
# affect density.

tracer_cpu = np.zeros((1, 6, n_elem))
tracer = tracer_equation.new_vector()
tracer.set(tracer_cpu)

temp_cpu = temp_cpu = temp0.copy().T[None, :, :]
temperature = temperature_equation.new_vector()
temperature.set(temp_cpu)

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

z_gpu = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv_gpu = problem.problemUV.new_vector()


# %%
# We prepare the output files for the standard exports. The 2D and 3D fields
# are written to XDMF.

path2d = "run2d"
path3d = "run3d"
mesh2d.create_xdmf(path2d)
mesh.create_xdmf(z0, path3d)


# %%
# Next, we choose the time-stepping parameters. We also compute CFL-based time
# step estimates to assess whether the chosen time step is reasonable.

iout = 0
t = 0
dt = 300
M = 20
ddt = dt / M
Tf = 40 * 3600 * 24
nstep = int(Tf / dt)
dt_max, ddt_max = problem.compute_cfl(u_max=u0, drho=2)

# %%
# Finally, we run the simulation. At regular intervals we export the 2D and 3D
# solution fields so that the evolution of the plume and tracer dome can be
# visualized afterwards.

problem.do_export(
    path2d, path3d, iout, t, z_gpu, uv_gpu, [temperature, tracer], rho=False
)

if mpi.rank == 0:
    print("End of init -- time = " + str(time.time() - tic) + "s")
    print("cfl : dt = %f, ddt = %f" % (dt_max, ddt_max))
    print("Simulation begins with M = %i, ddt = %f, dt = %f" % (M, ddt, dt))

tic = time.time()
for iiter in range(nstep):
    # We use the midpoint scheme here for the 3D update. Other time integrators
    # can be tested by replacing the line below with the appropriate iterator.
    # t = slim3d.iterate_impl(problem, t, dt, M, 1, z_gpu, uv_gpu, [temperature, tracer])
    t = slim3d.iterate_midpoint(
        problem, t, dt, M, 1, z_gpu, uv_gpu, [temperature, tracer]
    )

    # Export the solution approximately one hundred times during the run.
    if (iiter + 1) % (int(nstep / 1000)) == 0:
        iout += 1
        problem.do_export(
            path2d,
            path3d,
            iout,
            t,
            z_gpu,
            uv_gpu,
            [temperature, tracer],
            rho=False,
        )

        if mpi.rank == 0:
            print(
                f"Output {iout} at t = {t / 3600:.2f}h, {si_format((time.time() - tic) / (iiter + 1), 3)}s/iter"
            )

if mpi.rank == 0:
    print("finish simulation", time.time() - tic)