# %%
# 2D advection-diffusion testcase
# ===============================
# This example solves a simple shallow-water and tracer transport problem on a
# square domain. Two tracer components are transported by an imposed inflow and
# diffuse inside the domain.
#
# The setup is mainly intended as a compact testcase for tracer advection,
# diffusion, and open-boundary forcing.

import numpy as np
import slim4
from slim4 import slim, slim2d, tools


# %%
# First, we load the mesh. If the simulation is run in parallel, the mesh is
# partitioned automatically before being loaded.

g = 9.80616
meshfile = "square.msh"

if slim4.mpi.size > 1:
    partitioned_meshfile = slim2d.partition_mesh(meshfile)
    mesh = slim2d.Mesh(partitioned_meshfile)
else:
    mesh = slim2d.Mesh(meshfile)


# %%
# We define the bathymetry and create the shallow-water and tracer problems.
# The tracer problem carries two components.

nf = 200
bathymetry = slim.DataStatic.create_constant(mesh, [6])
shallow_water = slim2d.ShallowWater(mesh, g, bathymetry)
tracer_problem = slim2d.Tracer(mesh, shallow_water, bathymetry, g, nf=nf)
solution = tracer_problem.new_vector()
print(solution.size, solution.size < 2**31)

# %%
# We initialize the prognostic fields. The shallow-water state is prescribed in
# terms of free-surface elevation and transport, while both tracer components
# start from zero inside the domain.

transport_state = solution.get_sub(0).get()
tracer_state = solution.get_sub(1).get()

transport_state[0] = 0.0
transport_state[1] = 0.0
transport_state[2] = 2 * 6
tracer_state[0] = 0.0
tracer_state[1] = 0.0

solution.get_sub(0).set(transport_state)
solution.get_sub(1).set(tracer_state)


# %%
# Next, we set the tracer diffusion and source term.

# horizontal_diffusivity = slim.DataStatic.create_constant(mesh, 100)
tracer_source = slim.DataStatic.create_constant(mesh, 0.0)

# tracer_problem.set_diffusion(horizontal_diffusivity)
tracer_problem.set_source(tracer_source)


# %%
# We define the open-boundary data. The same tracer concentration, free-surface
# elevation, and transport are imposed on both open boundaries.

open_boundary_tracer = slim.DataStatic.create_constant(mesh, 1+np.arange(nf))
open_boundary_eta = slim.DataStatic.create_constant(mesh, 0)
open_boundary_transport = slim.DataStatic.create_constant(mesh, [0, 12])
open_boundaries = ["Open", "Open1"]

for boundary_name in open_boundaries:
    tracer_problem.set_open_boundary(
        boundary_name,
        tracer=open_boundary_tracer,
        eta=open_boundary_eta,
        u=open_boundary_transport,
        uv_is_transport=True,
        diffusion=False,
    )
    shallow_water.set_open_boundary(
        boundary_name,
        eta=open_boundary_eta,
        u=open_boundary_transport,
        uv_is_transport=True,
    )

tracer_problem.set_wall_boundary("Wall")
shallow_water.set_wall_boundary("Wall")


# %%
# We define a helper function to export the solution to XDMF. In addition to the
# solution fields, we monitor the maximum change in tracer concentration since
# the previous export.

output_name = "advection"
mesh.create_xdmf(output_name)


def export_solution(export_index, t, state, previous_tracer_state):
    # transport = state.get_sub(0).get()
    # tracer_components = state.get_sub(1).get()

    # mesh.append_to_xdmf(
    #     output_name,
    #     export_index,
    #     t,
    #     [
    #         ("H", transport[0]),
    #         ("Hu", transport[1:]),
    #         ("C1", tracer_components[0]),
    #         ("C2", tracer_components[1]),
    #     ],
    # )

    # max_difference = slim4.mpi.max_scalar(
    #     np.abs(previous_tracer_state - tracer_components).max()
    # )
    # if slim4.mpi.rank == 0:
    #     print(f"Export {export_index}, DIFF : {max_difference:.3e}")

    # previous_tracer_state[:] = tracer_components
    return export_index + 1


# %%
# Finally, we choose the time integration method and run the simulation. This
# example defaults to implicit Euler, but an explicit Runge-Kutta option is also
# kept for comparison.

implicit = False
dt = shallow_water.cfl
nsub = 100
t_end = 50000

if implicit:
    linear_system = tracer_problem.new_linear_system("")
    dt *= nsub

export_index = 0
t = 0
# previous_tracer_state = solution.get_sub(1).get().copy()
# export_index = export_solution(export_index, t, solution, previous_tracer_state)


from time import perf_counter as clock

ts = np.zeros(6)
iteration = 0
slim.synchronize()
ts[0] = clock()
while t < t_end:
    if implicit:
        t = slim.iterate_implicit_euler(tracer_problem, linear_system, solution, t, dt)
    else:
        t = slim.iterate_low_storage_runge_kutta(
            tracer_problem, solution, t, dt, 2/13, nsub
        )
    iteration += 1
    slim.synchronize()
    ts[iteration] = clock()
    # export_index = export_solution(export_index, t, solution, previous_tracer_state)
    print(iteration)
    if iteration == 5:
        break

dts = np.diff(ts)
print(dts)
print(np.mean(dts), np.std(dts), np.min(dts), np.max(dts))