# %%
# Stommel's gyre
# ==============
# This example shows how to run a simple 2D simulation with
# Stommel's box model of the atlantic gyre as an example
from slim4 import slim, slim2d, mpi
from slim4.slim import si_format
import numpy as np
import time

# %%
# First, let's load the mesh. If we are running in parallel, we need to partition the mesh first.

if mpi.size > 1:
    partname = slim2d.partition_mesh("square.msh")
    mesh = slim2d.Mesh(partname)
else:
    mesh = slim2d.Mesh("square.msh")
mesh.create_xdmf("stommel", cg=True)

# %%
# Now we can set the parameters of the problem. We can choose to run with either a linear drag (gamma) or a viscosity (kappa).

munk = True
if munk:
    kappa = 1e4
    gamma = 0.0
else:
    gamma = 1e-6
    kappa = 0.0
g = 9.80616
rho_0 = 1000
bath = slim.DataStatic.create_constant(mesh, 1000)


# %%
# Now we can create the problem and set forcings and the boundary conditions.
# We will use a simple sinusoidal surface stress to drive the flow, and a linear Coriolis parameter that varies with latitude.

problem = slim2d.ShallowWater(mesh, g, bath, rho_0, wave_only=False)
gamma_data = slim.DataStatic.create_constant(mesh, gamma)
problem.add_dissipation(gamma_data)

kappa_data = slim.DataStatic.create_constant(mesh, kappa)
problem.set_viscosity(kappa_data)

tau = slim.DataStatic(mesh, nfields=2)
tau.set(np.array([[0.1], [0]]) * [np.sin(np.pi * mesh.xnodes[1] / 1e6)])
problem.set_surface_stress(tau)

cor = slim.DataStatic(mesh, nfields=1)
cor.set(1e-4 + 2e-11 * mesh.xnodes[1])
problem.set_coriolis(cor)

problem.set_wall_boundary("North")
problem.set_wall_boundary("South")
problem.set_wall_boundary("West")
problem.set_wall_boundary("East")


# %%
# Next, we select the time-stepping scheme and the time step. We can choose to run with either an implicit or an explicit scheme.

implicit = False
tmax = 20000  # seconds
dt = 10000.0 if implicit else 2.0
substeps = 1 if implicit else 1000
nstep = int(tmax // (dt * substeps))


# %%
# We define a helper function to export the solution to XDMF files. We will export the free surface elevation (eta) and the velocity (u) fields. The velocity is computed from the transport fields by dividing by the total depth (eta + bathymetry).
hm = bath.eval_mesh()


def export(iout, t, solution):
    solution = solution.get()
    mesh.append_to_xdmf(
        "stommel",
        iout,
        t,
        [
            ("eta", solution[0]),
            ("u", solution[1:] / (solution[[0]] + hm)),
        ],
    )
    return iout + 1


# %%
# If we are running with an implicit scheme, we need to create a linear system to solve at each time step.
# We will use a simple LU factorization for this example, but more advanced solvers can be used for larger or distributed problems.
if implicit:
    sys = problem.new_linear_system("-pc_type lu")

# %%
# Finally, we can run the simulation and export the results to XDMF files that can be visualized with Paraview. We will also print some timing information to monitor the performance of the simulation.
iout = 0
t = 0.0
solution = problem.new_vector()
iout = export(iout, t, solution)

ts = np.zeros((nstep + 1))
tinit = time.time()
for i in range(nstep):
    if implicit:
        for i in range(substeps):
            t = slim.iterate_midpoint(problem, sys, solution, t, dt)
    else:
        t = slim.iterate_low_storage_runge_kutta(problem, solution, t, dt, 2, substeps)

    ts[i + 1] = time.time() - tinit
    periter = (ts[i + 1] - ts[i // 2]) / (i + 1 - i // 2) / substeps
    eta = ts[i + 1] + (nstep - (i - 1) * substeps) * periter
    iout = export(iout, t, solution)
    if mpi.rank == 0:
        print(
            f"iter {(i + 1) * substeps}/{nstep * substeps}, t={t}, cputime={ts[i + 1]:.3f}s, ETA={eta:.1}s, {si_format(periter, 3)}s/iter, fact={si_format(dt / periter, 5)}"
        )

dts = np.diff(ts) / substeps
if mpi.rank == 0:
    print()
    print(
        f"median={si_format(np.median(dts), 3)}s, mean={si_format(np.mean(dts), 3)}s, std={si_format(np.std(dts), 3)}s, min={si_format(np.min(dts), 3)}s"
    )
    print("factor with median time : %.2f" % (dt / np.median(dts)))