# %%
# Benchmark run
# -------------
#
# This script runs a short 3D SLIM benchmark on a partitioned idealized mesh and
# writes the timing results to benchmarks.sqlite.
#
# It is normally launched automatically by benchmark.sh rather than run
# directly.

import argparse
import sqlite3
import json
import socket
import numpy as np
from slim4 import mpi, scalar, slim, slim2d, slim3d, pre3d
from slim4.slim import si_format
from time import perf_counter as clock


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("-nz", dest="nz", type=int, help="An nz value")
    parser.add_argument("-m", dest="scale", type=str, help="The scale")
    parser.add_argument("-devname", dest="devname", type=str, help="Device name")
    args = parser.parse_args()
    return args.nz, args.scale, args.devname


# %%
# We parse the benchmark configuration and load the partitioned mesh.
nz, scale, devname = parse_arguments()
ftype = "fp32" if scalar == np.float32 else "fp64"
mesh2d = slim2d.Mesh(f"meshes/mesh{scale}_split/{mpi.size}/{mpi.rank}.msh")
n_tri_tot = mpi.sum_idx(mesh2d.n_triangles)
if mpi.rank == 0:
    print(f"    Ntri = {n_tri_tot}")

# %%
# We extrude the 2D mesh and initialize the 3D hydrodynamic problem.
bath2d = slim.DataStatic.create_constant(mesh2d, 1000)
mesh, z0, _bath2d = pre3d.extrude(
    mesh2d, bath2d, pre3d.VMapProportional(np.ones(nz, dtype=scalar))
)
X = mesh.create_coordinate_vector(z0)

g = 9.80616
rho0 = 1000
problem = slim3d.Problem(mesh, g, bath2d, rho0)
sw3d = problem.problemUV
ts_equation = problem.add_tracer(2, "temp_sal", limiter_type="past_bound")
n_elem = mesh.z_layers[-1]
problem.set_wall_boundary("Wall")

# %%
# We define idealized physical parameters and turbulence closures representative
# of realistic ocean-model configurations.
gamma = slim.DataStatic.create_constant(mesh2d, [0.01])
cor = slim.DataStatic.create_constant(mesh2d, 1.2e-4)
smago_coeff = slim.DataStatic.create_constant(mesh, value=0.02)
okubo_coeff = slim.DataStatic.create_constant(mesh, value=2e-4)

sw3d.set_dissipation(gamma, log_profile=False)
sw3d.set_coriolis(cor)
sw3d.set_horizontal_viscosity(smago_coeff, "smago")
for name, problem_tracer in problem.problem_tracer.items():
    problem_tracer.set_horizontal_diffusivity(okubo_coeff, "okubo")
problem.add_turbulence()

# %%
# The initial temperature field is initialized with a baroclinic front and a
# localized perturbation to trigger instabilities.
temp_sal = ts_equation.new_vector()
temperature = temp_sal.get_sub(0)
salinity = temp_sal.get_sub(1)
problem.set_state_equation_linear([(temperature, -0.2), (salinity, 0.1)])

z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = problem.problemUV.new_vector()

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_from_host(temp_init)
salinity.fill(35.0)


# %%
# We compute stable time steps and configure the benchmark duration.
dt3d, dt2d = problem.compute_cfl()
dt3d = min(dt3d, 10 * dt2d)  # realistic value for M : 10-15
dt = dt3d
ddt = dt2d
M = dt / ddt
t = 0.0
nstep = 100  # for benchmark

# %%
# Run the benchmark and record iteration timings.
ts = np.zeros(nstep + 1)
tic = clock()
for i in range(nstep):
    t = slim3d.iterate_midpoint(problem, t, dt, M, 1, z, uv, [temp_sal])
    slim.synchronize()
    ts[i + 1] = clock() - tic
    if (i + 1) > 10 and ts[i + 1] > 10:
        break
ts = ts[: i + 2]

if mpi.max_idx(uv.has_nan()):
    raise ValueError("NaN values in velocity")

# %%
# Timing statistics are written to a shared SQLite database for later analysis.
if mpi.rank == 0:
    dts = np.diff(ts[1:])  # skip first value to avoid cold start
    print(f"    min dt {si_format(dts.min(), 4, ' ')}s")
    print(f"    median dt {si_format(np.median(dts), 4, ' ')}s")
    print(f"    mean dt {si_format(dts.mean(), 4, ' ')}s")

    conn = sqlite3.connect("benchmarks.sqlite", timeout=60)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=60000")

    conn.execute("""
    CREATE TABLE IF NOT EXISTS benchmark (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        hostname TEXT NOT NULL,
        devname TEXT NOT NULL,
        dtype TEXT NOT NULL,
        n_triangles INTEGER NOT NULL,
        n_layers INTEGER NOT NULL,
        mpi_size INTEGER NOT NULL,
        timing_json TEXT NOT NULL,
        created_at TEXT DEFAULT CURRENT_TIMESTAMP
    )
    """)

    conn.execute(
        """
    INSERT INTO benchmark
    (hostname, devname, dtype, n_triangles, n_layers, mpi_size, timing_json)
    VALUES (?, ?, ?, ?, ?, ?, ?)
    """,
        (
            socket.gethostname(),
            devname,
            ftype,
            n_tri_tot,
            nz,
            mpi.size,
            json.dumps(dts.tolist()),
        ),
    )

    conn.commit()
    conn.close()