# %%
# Run the Great Barrier Reef simulation
# -------------------------------------
# This is the final step of the workflow. It loads the preprocessed mesh and
# forcing files produced by prepro.py, configures the 3D hydrodynamic and tracer
# problem, and advances the simulation in time.
#
# The model uses open-boundary elevation, velocity, temperature, and salinity;
# atmospheric forcing; reef-dependent bottom roughness; turbulence closure;
# Smagorinsky horizontal viscosity; Okubo tracer diffusivity; and a realistic
# temperature-salinity equation of state.

# %%
from time import perf_counter as clock
import numpy as np
import slim4
from slim4 import slim, slim3d, mpi
from slim4.tools import xdmf_abin_shrink, parse_time, format_time, dhms
import params as param


def log(*args, **kwargs) -> None:
    if mpi.rank > 0:
        return
    print(*args, **kwargs)


# %%
# We gather MPI information, define the preprocessing directory, and parse the
# initial time used for ramps and initial conditions.

comm_rank = mpi.rank
comm_size = mpi.size
prepro_dir = param.prepro_data_dir + "/"

# Time interval
tstart, tend = param.initial_time, param.final_time
t_init = parse_time(tstart)

# %%
# We load the partitioned 3D mesh and bathymetry created by prepro.py.
log("Loading mesh file...")
ishybrid = param.n_sigma > 0 and (
    (isinstance(param.z_levels, int) and param.z_levels > 0)
    or (isinstance(param.z_levels, list) and len(param.z_levels) != 0)
)
mesh, z0, bath2d = slim3d.Mesh.load(
    prepro_dir + "mesh3d_split", n_moving_max=param.n_moving_max
)
mesh2d = mesh.mesh2d
nt = mesh2d.x.shape[-1]

# %%
# We create the 3D problem and add the combined temperature-salinity tracer.
# The two tracer components are later exposed as named subvectors for output and
# for the equation of state.

# Initialize problems
g = 9.80616
rho_0 = 1000
n_elem = mesh.z_layers[-1]
run = slim3d.Problem(mesh, g, bath2d, rho_0)
problem_ts = run.add_tracer(2, "temp_sal", limiter_type="past_bound")
sw2d = run.problemUV2d
sw3d = run.problemUV

# %%
# We load time-dependent open-boundary data. The 3D velocity is converted to a
# transport by multiplying by the moving-layer thickness associated with the
# boundary elevation and bathymetry. A ramp is applied at the start of the run.

# Load data for boundary conditions
DT = param.ramp
u_open = slim.DataFile(prepro_dir + "slim4_uv", "u", mesh)
u_init = slim.DataFile(prepro_dir + "slim4_uv_init", "u", mesh)
eta_open = slim.DataFile(prepro_dir + "slim4_eta", "h", mesh2d)
hu_open = slim3d.DataHMultiply(u_open, eta_open, bath2d, z0)
hu_open.set_ramp(t_init, DT)
eta_open.set_ramp(t_init, DT)
ts_open = slim.DataFile(prepro_dir + "slim4_temp_sal_open", "temp_sal", mesh)
ts_init = slim.DataFile(prepro_dir + "slim4_temp_sal_init", "temp_sal", mesh)

# %%
# We apply wall boundaries and open boundaries. At open boundaries, elevation,
# transport, temperature, and salinity are all prescribed from the preprocessed
# files.

# Set boundary conditions
for tag in param.closed_tags:
    run.set_wall_boundary(tag)
for tag in param.open_tags:
    run.set_open_boundary(
        tag, eta=eta_open, u=hu_open, t={"temp_sal": ts_open}, uv_is_transport=True
    )

# %%
# We initialize the prognostic fields: moving vertical coordinates, 3D velocity,
# and temperature-salinity tracers. The initial temperature and salinity fields
# are read from the preprocessed ocean product.

# Create and initialize variables
z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = sw3d.new_vector()
ts_tracers = problem_ts.new_vector()
temperature = ts_tracers.get_sub(0)
temperature.set_name("temperature")
salinity = ts_tracers.get_sub(1)
salinity.set_name("salinity")

# Set initial conditions for temperature and salinity
ts_tracers.set(ts_init.eval_mesh(t_init))

# %%
# We apply wind stress, and optionally atmospheric pressure, from the atmospheric
# forcing files.

# Wind and atmospheric pressure
tau = slim.DataFile(prepro_dir + "slim4_wind", "tau", mesh2d)
wind = slim.DataFile(prepro_dir + "slim4_wind", "wind", mesh2d)
atm = slim.DataFile(prepro_dir + "slim4_wind", "psl", mesh2d)
tau.set_ramp(t_init, DT)
atm.set_ramp(t_init, DT)
run.set_surface_stress(tau=tau)
sw2d.set_atmospheric_pressure(atm)

# %%
# We apply surface and penetrating heat fluxes when requested. Additional
# temperature-reaction terms correct the fluxes using sea-surface temperature,
# air temperature, humidity, and latent/sensible heat components.

# Load data
surf_heat_open = slim.DataFile(prepro_dir + "slim4_atmo", "surf_flux", mesh2d)
pen_heat_open = slim.DataFile(prepro_dir + "slim4_atmo", "pen_flux", mesh2d)
pen_dist = slim.DataFile(prepro_dir + "slim4_atmo", "pen_dist", mesh2d)

# Surface water and heat fluxes
surf_heat_open.set_ramp(t_init, DT)
pen_heat_open.set_ramp(t_init, DT)

# Set surface heat fluxes
run.set_surface_flux("temp_sal", surf_heat_open)
run.set_penetrating_flux(pen_heat_open, pen_dist, temperature, z)

# Add corrections to the surface heat fluxes :
sst = slim.DataFile(prepro_dir + "slim4_atmo", "sst", mesh2d)
airtemp = slim.DataFile(prepro_dir + "slim4_atmo", "airtemp", mesh2d)
air_mix_ratio = slim.DataFile(prepro_dir + "slim4_atmo", "air_mix_ratio", mesh2d)
sat_mix_ratio = slim.DataFile(prepro_dir + "slim4_atmo", "sat_mix_ratio", mesh2d)
q_lat = slim.DataFile(prepro_dir + "slim4_atmo", "q_lat_wm2", mesh2d)
q_sens = slim.DataFile(prepro_dir + "slim4_atmo", "q_sens_wm2", mesh2d)
reaction_fluxes = slim3d.ReactionTemperatureCorrection(
    mesh,
    sst,
    air_mix_ratio,
    sat_mix_ratio,
    q_lat,
    q_sens,
    airtemp,
    wind,
    z,
    temperature,
)
run.add_reaction(reaction_fluxes)

# %%
# Static and parameterized physical terms: Coriolis, reef-dependent bottom
# roughness, turbulence, horizontal viscosity, tracer diffusivity, the equation
# of state, and velocity clipping.

# Coriolis
cor = slim4.slim.DataFile(prepro_dir + "cor", "cor", mesh2d)
run.problemUV.set_coriolis(cor)

# Disipation
reef_density = slim.DataFile(
    prepro_dir + "reef_density", "reef_density", mesh2d
).eval_mesh()
gamma = slim.DataStatic(mesh2d, "DG", 1)
base_z0b = 0.003  # roughness length for sandy bottom
hard_z0b = 0.2  # roughness length for hard bottom
gamma.set(base_z0b + reef_density * (hard_z0b - base_z0b))
sw3d.set_dissipation(gamma, log_profile=True)

# Turbulence closure for vertical viscosity and diffusivity
run.add_turbulence(z0b=gamma)

# Horizontal viscosity (SMAGO) and diffusivity (OKUBO)
smago_coeff = slim.DataStatic.create_constant(mesh, value=0.02)
okubo_coeff = slim.DataStatic.create_constant(mesh, value=2e-4)
sw3d.set_horizontal_viscosity(smago_coeff, "smago")
for name, problem_tracer in run.problem_tracer.items():
    problem_tracer.set_horizontal_diffusivity(okubo_coeff, "okubo")

# State equation for the density
run.set_state_equation_jackett(temperature, salinity)

# Clip velocity
max_uv = 3.0  # Max horizontal speed
realbath = slim.DataFile(prepro_dir + "bath", "realbath", mesh2d)
run.set_clipping_velocity(
    max_uv,
    slim.ScalarArrayDevice.create_zeros([1, 1, n_elem]),
    realbath=realbath,
)

run.set_clipping_velocity(
    3.0, slim.ScalarArrayDevice.create_zeros([1, 1, n_elem]), realbath=realbath
)

# %%
# We estimate stable 3D and 2D time steps from the CFL condition, then choose a
# model time step that divides the requested export interval exactly.

# Compute timesteps
rho = run.get_gpu_tmpvars().rho
run.problemUV.update_density(z, rho)
dt3d, dt2d = run.compute_cfl(u_max=max_uv, drho=rho.get())
dt3d *= param.dt_mult
nsub = int(np.ceil(param.export_dt / dt3d))
dt = param.export_dt / nsub
M = dt / dt2d
log(f"dt3d = {dt:.3f}, M = {M:.2f}")

# %%
# We initialize the temporal loop. In restart mode, previous save files are read
# and XDMF/abin outputs are shrunk so new results append cleanly.

# Initialize values for the temporal loop and load previous data if restarting
checkdt = 24
if param.restart:
    t, iout, _ = run.load(
        param.output_directory + "/saves", z, uv, [temperature, salinity]
    )
    xdmf_abin_shrink(
        param.output_directory + "/slim3D", divmod(iout, checkdt)[0] * checkdt + 1
    )
    xdmf_abin_shrink(
        param.output_directory + "/slim2D", divmod(iout, checkdt)[0] * checkdt + 1
    )
else:
    t = t_init
    iout = 0
iiter = 0

# %%
# We prepare the XDMF output files and write the initial state unless this is a
# restart. The export options include density and velocity-clipping diagnostics.

# Create xdmf files and define exports
path3d = param.output_directory + "/slim3D"
path2d = param.output_directory + "/slim2D"
export_kwargs = dict(
    Huv=False,
    turbulence=False,
    rho=True,
    clip_loc=run.uv_clip_loc,
    nitperex=nsub,
)
if not param.restart:
    mesh.create_xdmf(z0, path3d, cg=True)
    mesh2d.create_xdmf(path2d, cg=True)
    run.do_export(
        path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
    )

mpi.barrier()
log("starting iterations")

# %%
# We run the simulation. At each export interval, fields are written to disk,
# basic diagnostics are printed, and periodic save files are created for restart.

# Temporal loop
tic = toc = clock()
log()
while t + dt <= parse_time(tend):
    t = slim3d.iterate_impl(run, t, dt, M, 1, z, uv, [ts_tracers])
    # t = slim3d.iterate_midpoint(run, t, dt, M, 1, z, uv, [ts_tracers])

    periter = (clock() - tic) / (iiter + 1)
    if mpi.size == 1:
        log(
            f"\033[FIter {iiter + 1:d}, Time {format_time(t)}, Clock Time {dhms(clock() - tic)}, {periter:.3f} s per iter ({dt / periter:.1f}x)"
        )

    if ((iiter + 1) % nsub) == 0 or mpi.max_idx(uv.has_nan()):
        iout += 1
        run.do_export(
            path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
        )
        if max_uv is not None:
            run.uv_clip_loc.memset(0)

        uvmax = mpi.max_scalar(np.linalg.norm(uv.get(), axis=0).max())
        w_max = mpi.max_scalar(np.abs(run.get_gpu_tmpvars().w_gpu.get()).max())
        temp_cpu = temperature.get()
        sali_cpu = salinity.get()
        temp_min = mpi.min_scalar(temp_cpu.min())
        temp_max = mpi.max_scalar(temp_cpu.max())
        sali_min = mpi.min_scalar(sali_cpu.min())
        sali_max = mpi.max_scalar(sali_cpu.max())
        log(f"Output {iout}, Time {format_time(t)}, Clock time {dhms(clock() - tic)}")
        log(f"max uv = {uvmax:.2f}, max w = {w_max:.2f}")
        log(f"{temp_min:.3f} < temperature < {temp_max:.3f}")
        log(f"{sali_min:.3f} < salinity < {sali_max:.3f}")
        toc = clock()
        if iout % checkdt == 0:
            run.save(
                t,
                iout,
                iiter,
                param.output_directory + "/saves",
                z,
                uv,
                [temperature, salinity],
            )
    iiter += 1

log("Hydro done =)")

# %%