# %%
# Baroclinic instability test case
# ================================
# This example sets up a stratified configuration with a horizontally varying
# temperature front that triggers baroclinic instabilities. The script builds and
# partitions the mesh, initializes the temperature field, runs the simulation,
# and exports both standard outputs and optional online visualizations.

import os
from time import perf_counter as clock

import numpy as np
from matplotlib import pyplot as plt
from osgeo import gdal

from slim4 import mpi, slim, slim2d, slim3d, pre3d
from slim4.post import Rasterizer
from slim4.slim import si_format

gdal.UseExceptions()

# %%
# We first create and partition the mesh. The 2D mesh is extruded vertically
# with sigma layers, saved to disk, then repartitioned for the current number
# of MPI ranks before being loaded for the actual simulation.

tic = clock()
comm_rank = mpi.rank

mesh2d0 = slim2d.Mesh(slim2d.partition_mesh("mesh3dSlim.msh", "tmp/mesh-${rank}.msh"))
bath2d0 = slim.DataStatic.create_constant(mesh2d0, 1000)
nz = 20
mesh3d0, z_seed, _ = pre3d.extrude(mesh2d0, bath2d0, pre3d.VMapProportional(np.ones(nz)))
mesh3d0.save_mesh(z_seed, "mesh3d")

slim3d.partition_mesh("mesh3dSlim.msh", "mesh3d", "mesh3d_split")

mesh, z0, bath2d = slim3d.Mesh.load("mesh3d_split", n_moving_max=3)
mesh2d = mesh.mesh2d
X = mesh.create_coordinate_vector(z0)


# %%
# Next, we create the 3D problem and configure the momentum and tracer
# equations. We use a single tracer for temperature and apply no-slip walls on
# the lateral boundary.

g = 9.80616
rho0 = 1000
problem = slim3d.Problem(mesh, g, bath2d, rho0)
temperature_equation = problem.add_tracer(1, "temperature", limiter_type="past_bound")
n_elem = mesh.z_layers[-1]

problem.set_wall_boundary("Wall")


# %%
# We now define the physical parameters of the flow. The setup includes bottom
# friction, Coriolis forcing, and horizontal and vertical viscosity.

gamma = slim.DataStatic.create_constant(mesh2d, [0.01])
problem.problemUV.set_dissipation(gamma, log_profile=False)

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

nu_h = slim.DataStatic.create_constant(mesh, 20)
problem.problemUV.set_horizontal_viscosity(nu_h)

nu_v = slim.ScalarArrayDevice.create_from_host(1e-4 * np.ones([1, 6, n_elem]))
problem.problemUV.set_vertical_viscosity(nu_v)


# %%
# Optinally, the vertical viscosity and diffusivity can be derived from a vertical turbulence closure.
# This replaces the constant nu_v defined above and adds a turbulent diffusivity for the temperature tracer.

# problem.add_turbulence()


# %%
# We initialize the prognostic fields and auxiliary arrays used by the 3D
# model: the temperature tracer, the moving vertical coordinates, and the
# 3D velocity field. The temperature tracer is coupled to a linear equation
# of state so that temperature variations induce density anomalies.

temperature = temperature_equation.new_vector()
problem.set_state_equation_linear([(temperature, -0.2)])
z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = problem.problemUV.new_vector()


# %%
# The initial temperature field is built from a vertical background profile and
# a laterally varying front. A localized perturbation is added to seed the
# instability and accelerate the onset of baroclinic structures.

temp_bot = 10.1
temp_top = 13.1
z_bot = -975
d_temp = 1.2
y_0 = 250e3
y_a = 40e3
k = 3
Lx = 160e3
d_y = 40e3

x, y, _ = X
temp_background = temp_bot + ((temp_top - temp_bot) * (z_bot - z0[0]) / z_bot)
front_center = y_0 - y_a * np.sin(2 * np.pi * k * x / Lx)
perturbation_center = y_0 - y_a / 2 * np.sin(np.pi * ((x - 110e3) / 20e3))
temp_perturbation = 0.3 * (1 - ((y - perturbation_center) / (d_y / 2)))

temp_init = np.zeros((6, n_elem))

mask_hot = y > front_center + d_y
temp_init[mask_hot] = temp_background[mask_hot]

mask_cold = y < front_center
temp_init[mask_cold] = (temp_background - d_temp)[mask_cold]

mask_transition = ~mask_hot & ~mask_cold
temp_init[mask_transition] = (
    temp_background - d_temp * (1 - ((y - front_center) / d_y))
)[mask_transition]

mask_perturbation = (
    (x > 110e3)
    & (x < 130e3)
    & (y > perturbation_center - (d_y / 2))
    & (y < perturbation_center + (d_y / 2))
)
temp_init[mask_perturbation] += temp_perturbation[mask_perturbation]

temp_init = temp_init.reshape((1, 6, -1))

temperature.set(temp_init)

# %%
# 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, cg=False)
mesh.create_xdmf(z0, path3d, cg=False)

# %%
# Optionally, we export time-averaged 2D and 3D transports. These are required
# for offline tracer computations

path2d_transport = "transport_2d"
path3d_transport = "transport_3d"

# %%
# We then choose the time step from the CFL condition. The simulation is run
# for 200 days, with exports evenly spaced in time.

t = 0.0
dt3d, dt2d = problem.compute_cfl()
tmax = 200 * 24 * 3600
export_dt = tmax / 400
nsub = int(np.ceil(export_dt / dt3d))
dt = export_dt / nsub
nstep = int(np.ceil(tmax // dt))
M = dt / dt2d
ddt = dt / M

# %%
# The online visualization is optional. We keep its setup separate from the
# main time loop so that the simulation logic stays easier to read. This is
# provided as an example: here we visualize the vorticity, but the same
# approach can be used for any field.


def setup_online_visualization(mesh, mesh2d):
    os.makedirs("img", exist_ok=True)
    raster = Rasterizer(mesh)
    xmin, xmax, ymin, ymax = (
        mpi.min_scalar(mesh2d.x[0].min()),
        mpi.max_scalar(mesh2d.x[0].max()),
        mpi.min_scalar(mesh2d.x[1].min()),
        mpi.max_scalar(mesh2d.x[1].max()),
    )
    raster.set_bbox(xmin, xmax, ymin, ymax, height=1000, width=-1)
    return raster, "seismic"


def export_online_visualization(mesh, raster, cmap, z, uv, iout):
    vorticity = mesh.vorticity(z, uv)
    surf_vort = raster.render(vorticity, render_type="surface")

    if mpi.rank != 0:
        return

    surf_vort_cpu = surf_vort.get()

    plt.figure(figsize=(3, 8))
    plt.imshow(surf_vort_cpu[0], cmap=cmap, vmin=-5e-4, vmax=5e-4)
    plt.xticks([])
    plt.yticks([])
    cbar = plt.colorbar(orientation="horizontal", fraction=0.046, pad=0.04)
    cbar.set_ticks([-5e-4, 0, 5e-4])
    plt.tight_layout()
    plt.savefig(f"img/{iout}.png", dpi=300, bbox_inches="tight")
    plt.close()

    bands, nrows, ncols = surf_vort_cpu.shape
    driver = gdal.GetDriverByName("GTiff")
    options = ["COMPRESS=ZSTD"]  # fast lossless compression
    # options = ["COMPRESS=LERC_ZSTD", "MAX_Z_ERROR=1e-4"]  # lossy export

    ds = driver.Create(
        f"img/{iout}.tif",
        ncols,
        nrows,
        bands,
        gdal.GDT_Float32,
        options=options,
    )
    ds.SetGeoTransform(raster.geotransform)
    # ds.SetProjection(tools.parse_projection(mesh.mesh2d.projection).ExportToWkt())

    for iband in range(bands):
        band = ds.GetRasterBand(iband + 1)
        band.WriteArray(surf_vort_cpu[iband].astype(np.float32))
        band.SetNoDataValue(np.nan)

    ds.FlushCache()
    ds = None


online_vis = True
if online_vis:
    raster, cmap = setup_online_visualization(mesh, mesh2d)


# %%
# Finally, we export the initial state and advance the model in time. The 3D
# system is integrated with the midpoint scheme, while alternative time
# integrators are left in comments for convenience.

iout = 0
iout_t = 0

iout_t = problem.export_mean_transport(
    path2d_transport, path3d_transport, iout_t, t, z, uv
)
iout = problem.do_export(path2d, path3d, iout, t, z, uv, [temperature], rho=True)

if comm_rank == 0:
    print(f"End of init, took {clock() - tic:.2f}s")
    print(
        f"Simulation begins with M = {M:.3f}, dt = {si_format(dt, 4, '')}s, ddt = {si_format(ddt, 3, '')}s"
    )

tic = clock()
ts = np.zeros(nsub + 1)
ts[0] = clock()

for i in range(nstep):
    # Options for the 3D scheme:
    # t = slim3d.iterate_impl(problem, t, dt, M, 1, z, uv, [temperature])
    # t = slim3d.iterate_expl(problem, t, dt, M, 1, z, uv, [temperature])
    # t = slim3d.iterate_impl_rk1(problem, t, dt, M, 1, z, uv, [temperature])
    # t = slim3d.iterate_expl_rk1(problem, t, dt, M, 1, z, uv, [temperature])
    t = slim3d.iterate_midpoint(problem, t, dt, M, 1, z, uv, [temperature])

    # Options for the 2D scheme:
    # solver2d = slim3d.TemporalIterator.rkls(2 / 13)
    # solver2d = slim3d.TemporalIterator.dirk(2)
    # t = slim3d.iterate_impl(problem, t, dt, M, 1, z, uv, [temperature], solver2d)

    ts[(i + 1) % (nsub + 1)] = clock()

    if comm_rank == 0:
        periter = si_format((clock() - tic) / (i + 1), 3)
        print(f"\riter {i + 1}, time {t:.2f}, {periter}s/iter", end="")

    if (i + 1) % nsub == 0:
        dts = np.diff(ts)
        print()
        med_time = np.median(dts)
        print(f"median time per iter: {1000 * med_time:.2f} ms")

        iout_t = problem.export_mean_transport(
            path2d_transport, path3d_transport, iout_t, t, z, uv
        )
        iout = problem.do_export(
            path2d, path3d, iout, t, z, uv, [temperature], rho=True
        )

        if online_vis:
            export_online_visualization(mesh, raster, cmap, z, uv, iout)

        if comm_rank == 0:
            print(f"    Export {iout}, time {t:.2f}, since {(clock() - tic):.2f} s")

        ts[0] = clock()

if comm_rank == 0:
    print("finish simulation", clock() - tic)