# %%
# Preprocess data for SLIM4
# -------------------------
# This is the fourth step of the Great Barrier Reef workflow. It converts the
# static and time-varying datasets prepared by the previous scripts into SLIM4
# data files on the model mesh.
#
# The script builds the 3D mesh and bathymetry, computes static fields such as
# Coriolis and reef density, and interpolates atmospheric, tidal, ocean-current,
# temperature, and salinity products onto the grids and boundaries needed by
# run.py.

import os
import numpy as np
from slim4 import (
    slim,
    slim2d,
    slim3d,
    pre3d,
    tpxo,
    mpi,
    particles,
)
from slim4.tools import parse_time, VectorRotation, read_polygons_shp
from slim4.netcdf_to_abin import nc2abin
from slim4.pre import Struct, clip_inclusive
from gbrtools import bathy_on_mesh

import params as param


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


def smith_banke(uv_wind: np.ndarray):
    """Compute wind stress with the Smith and Banke drag law.

    Parameters
    ----------
    uv_wind : ndarray of shape (2, n)
        Wind velocity components in m/s.

    Returns
    -------
    ndarray of shape (2, n)
        Wind stress components.
    """
    atm_density = 1.25
    f0, f1, f2 = 0, 6.3e-4, 6.6e-5
    norm = np.hypot(*uv_wind)
    coeff = f0 + f1 * norm + f2 * norm**2
    return atm_density * coeff[None, :] * uv_wind


tstart = param.initial_time
tend = param.final_time

timelaps = np.array([parse_time(tstart), parse_time(tend)])
time_tides = np.arange(timelaps[0], timelaps[-1] + 3599, 3600)
time_check = np.arange(timelaps[0], timelaps[-1], param.check_dt)

base_dir = param.nc_data_dir + "/"
prepro_dir = param.prepro_data_dir + "/"
if mpi.rank == 0:
    os.makedirs(prepro_dir + "check", exist_ok=True)
mpi.barrier()


def get_files(path, vars):
    path = os.path.join(base_dir, path)
    r = []
    for var in vars:
        ncfiles = os.listdir(path)
        ncfiles = [f for f in ncfiles if f.startswith(var)]
        ncfiles = [os.path.join(path, f) for f in ncfiles if f.endswith(".nc")]
        for f in ncfiles:
            if not os.path.exists(f[:-3]):
                print(f"Converting {f[:-3]} to abin format...")
                nc2abin(f)
        ncfiles = [f[:-3] for f in ncfiles]
        ncfiles.sort()
        if len(ncfiles) == 0:
            raise FileNotFoundError(f"No file found for variable {var} in path {path}")
        r.append(ncfiles)

    return r


def get_struct(path, vars):
    """Load 2D/3D forcing variables as a merged SLIM4 Struct.

    Parameters
    ----------
    path : str
        Subdirectory of ``base_dir`` containing the files.
    vars : list of str
        Variables to load.

    Returns
    -------
    Struct
        Merged structure containing the requested variables.
    """
    flists = get_files(path, vars)
    structs = [Struct.from_files(flist, [var]) for (flist, var) in zip(flists, vars)]
    return Struct.merge_structs(structs)


# %%
# For the bathymetry source used for preprocessing, the preferred approach is
# to provide a list of rasters, which are loaded in order of priority
# (earlier files take precedence where they overlap). This preserves the
# highest available resolution locally.
#
# Single merged rasters are convenient for prototyping but less accurate.
tif = [
    "data/bathymetry/" + s
    for s in [
        "Great_Barrier_Reef_A_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_B_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_C_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_D_2020_30m_MSL_cog.tif",
        "Australian_Bathymetry_Topography__Torres_Strait__2023_30m_MSL_cog.tif",
        "Great_Barrier_Reef_2020_100m_MSL_cog.tif",
        "GulfOfPapua_100m_bathymetry.tif",
    ]
]

# Alternative (faster, lower fidelity):
# tif = "data/bathymetry/merged/merge_d100.tif"   # good compromise
# tif = "data/bathymetry/merged/merge_d1000.tif"  # very fast prototyping


# %%
# Next, we read and (temporarily) partition the horizontal mesh generated by mesh.py.
# This 2D partition is only used during preprocessing. After extrusion, the mesh
# is repartitioned in 3D so that both the horizontal triangles and vertical
# prisms are balanced across MPI ranks.
#
# Bathymetry is interpolated onto the mesh, clipped to this minimum depth, and
# smoothed in shallow regions. The mesh is then extruded using a hybrid vertical
# coordinate. Because the computational cost depends on the number of vertical
# layers, the final partitioning is performed on the 3D mesh to properly balance
# the workload.
def mesh_and_bathy():
    mesh_file = param.mesh_file
    if mpi.size > 1:
        pfile = prepro_dir + "/mesh/${rank}/mesh.msh"
        mesh_file = slim2d.partition_mesh(param.mesh_file, pfile=pfile)
    mesh2d0 = slim2d.Mesh(mesh_file)

    lon0, lat0 = slim.proj_transform(mesh2d0.projection, "lonlat", mesh2d0.x)[
        [[0], [1]]
    ]
    tides0 = tpxo.GridTPXO(
        mesh2d0, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone", npass=30
    )
    tidal_amplitude0 = (
        tides0.tides["h"]
        .approx_amplitude_at_points(
            np.array([lon0.ravel(), lat0.ravel()]).T, timelaps[-1]
        )
        .reshape(lon0.shape)
    )
    min_depth0 = tidal_amplitude0 + param.min_depth

    bath_dg0 = bathy_on_mesh(mesh2d0, tif, nodal=False, cg=True, rule_depth=2)
    bath_dg0 = np.maximum(bath_dg0, min_depth0)
    deep_areas = bath_dg0 > param.depth_sigma
    bath_deep = bath_dg0[deep_areas]
    bath_dg0[deep_areas] = param.depth_sigma

    bath_dg0 = mesh2d0.smooth_bath(bath_dg0, coefficient=0.4, minimum_delta=0.3)
    # bath_dg0 = mesh2d0.smooth_bath_haney(bath_dg0, coefficient=1.5, minimum_delta=0.3)
    bath_dg0[deep_areas] = bath_deep
    bath0 = slim.DataStatic(mesh2d0, "DG", 1, value=bath_dg0)

    shallow_vmap = pre3d.VMapSongHaidvogel(param.n_sigma, 4, 0.5, 5)
    vmap = pre3d.VMapHybrid.create_exponential_deep(
        param.depth_sigma, shallow_vmap, param.z_levels, True, bath0
    )
    mesh0, z_, _ = pre3d.extrude(mesh2d0, bath0, vmap)
    mesh0.save_mesh(z_, prepro_dir + "mesh3d")
    slim3d.partition_mesh(
        param.mesh_file, prepro_dir + "mesh3d", prepro_dir + "mesh3d_split"
    )


mesh_and_bathy()

# %%
# We load the repartitioned 3D mesh and recompute the real bathymetry on the
# final 2D mesh. This field is kept separately from the smoothed model
# bathymetry and is later used for diagnostics and velocity clipping.

mesh, z, bath = slim3d.Mesh.load(
    prepro_dir + "mesh3d_split", n_moving_max=param.n_moving_max
)
mesh2d = mesh.mesh2d
realbath = bathy_on_mesh(mesh2d, tif, nodal=False, cg=True, rule_depth=2)
realbath = np.maximum(realbath, -1.0)
deep_areas = realbath > param.depth_sigma
bath_deep = realbath[deep_areas]
realbath[deep_areas] = param.depth_sigma
realbath = mesh2d.smooth_bath(realbath, coefficient=2.0, minimum_delta=0.3)
realbath[deep_areas] = bath_deep
slim.DataFile.write(
    prepro_dir + "bath", [("bath", bath.eval_mesh()), ("realbath", realbath)], "DG"
)
mesh2d.create_xdmf(param.output_directory + "/bath", cg=False)
mesh2d.append_to_xdmf(
    param.output_directory + "/bath",
    0,
    0.0,
    [("bath", bath.eval_mesh()), ("realbath", realbath)],
)
lon, lat = slim.proj_transform(mesh2d.projection, "lonlat", mesh2d.x)[[[0], [1]]]

bounds = lon.min(), lon.max(), lat.min(), lat.max()
mpi.barrier()


# %%
# We prepare static 2D fields used by the simulation: the Coriolis parameter and
# reef-density-based bottom roughness.
def prepro_coriolis():
    """Compute and write the Coriolis parameter on the 2D mesh."""
    cor = 2 * (2 * np.pi / 86400) * np.sin(np.deg2rad(lat))
    slim.DataFile.write(prepro_dir + "cor", [("cor", cor)], "DG")
    mpi.barrier()
    log("coriolis done")


def prepro_reef_denisty():
    """Compute reef coverage density on mesh elements and write it to disk."""
    reef_shp = "data/reefs.gpkg"
    polymesh = particles.PolygonMeshIntersection(
        mesh2d, *read_polygons_shp(mesh2d, reef_shp)
    )
    reef_density = polymesh.get_density().reshape((1, 3, -1))
    reef_density = np.clip(reef_density, 0, 1)
    reef_density_cg = mesh2d.gather_dg_to_cg("mean_area", reef_density)
    sreef_density = reef_density_cg[0][mesh2d.element_nodes][:, :, None].T

    slim.DataFile.write(
        prepro_dir + "reef_density", [("reef_density", sreef_density)], "DG"
    )
    mesh2d.create_xdmf(param.output_directory + "/reef_density")
    mesh2d.append_to_xdmf(
        param.output_directory + "/reef_density",
        0,
        0,
        [("reef_density", sreef_density)],
    )
    mpi.barrier()
    log("Reef density done")


# %%
# Next, we prepare the 2D time-dependent surface and open-boundary forcings, including
# wind surface stress, heat fluxes, precipitation and sea level
def prepro_wind():
    """Interpolate wind, rotate it to mesh coordinates, and compute wind stress."""
    log("Processing WIND velocity...")
    wind_vars = ["uas", "vas"]
    wind_vars.append("psl")  # if use atmospheric pressure

    wind = get_struct("wind", wind_vars)
    csr = wind.compute_csr_cg_2d(mesh2d)
    rot_wind = VectorRotation(mesh2d, wind.get_lonlat_grid(csr))
    with slim.DataFileWriter(
        prepro_dir + "slim4_wind", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(wind.get_time(*timelaps)):
            [(_, windx), (_, windy)] = wind.get_data(csr, t, ["uas", "vas"])
            windx, windy = rot_wind.rotate(windx, windy)
            ws = np.vstack([windx, windy])
            tau = smith_banke(ws)
            wind_data = [("wind", ws), ("tau", tau)]
            wind_data += wind.get_data(csr, t, ["psl"])
            writer.append(wind_data, t)

    mpi.barrier()


def prepro_elevation():
    """Combine ocean-model sea level with tidal elevation on open boundaries."""
    log("Processing Sea Surface Elevation...")
    tides = tpxo.GridTPXO(
        mesh2d, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone"
    )
    eta = get_struct("hydro2d", ["eta_t"])

    # tides typically have a finer grid, so interpolate eta on the tide grid
    eta.set_reinterp_coordinates(tides.lon, tides.lat)
    csr = eta.compute_csr_cg_2d(mesh2d, param.open_tags)
    tides.set_points(eta.get_lonlat_grid(csr))
    with slim.DataFileWriter(
        prepro_dir + "slim4_eta", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_tides):
            [(_, h)] = eta.get_data(csr, t, ["eta_t"])
            h += tides.get_eta(t)
            writer.append([("h", h)], t)

    mpi.barrier()


def rel_hu_from_dew(dew, tair):
    """Calculate relative humidity from dew point and air temperature.

    Parameters
    ----------
        dew : NDarray
            Dewpoint temperature in °C.
        tair : NDarray
            Air temperature in °C.
    Returns
    -------
        rel_hum : NDarray
            Relative humidity as a ratio between 0 and 1.
    """
    b = 17.625
    c = 243.04  # degree Celsius
    return np.exp(b * c * (dew - tair) / ((c + tair) * (c + dew)))


def dewpoint(rel_hum, temp):
    """
    Parameters:
        rel_hum -- relative humidity as a ratio between 0 and 1.
        temp -- air temperature in degree Celsius
    """
    b = 17.625
    c = 243.04  # degree Celsius
    gamma = (b * temp) / (c + temp) + np.log(rel_hum)
    dew_point = (c * gamma) / (b - gamma)
    return dew_point


def vapor_pressure(temp):
    """Calculate saturation vapor pressure from temperature.

    Parameters
    ----------
        temp : NDarray
            Temperature in °C.
    Returns
    -------
        vapor_press : NDarray
            Saturation vapor pressure in Pa.
    """
    return 6.133523 * np.exp(17.502 * temp / (240.97 + temp))


def prepro_atmo():
    """Prepare atmospheric heat-flux and humidity fields for the model."""
    log("Processing atmospheric data...")
    atmo_vars1 = [
        "hfls",  # surface upward latent heat flux [W/m²]
        "hfss",  # surface upward sensible heat flux [W/m²]
        "rsds",  # surface downwelling shortwave radiation [W/m²]
        "rlus",  # surface upwelling longwave radiation [W/m²]
    ]
    atmo_vars2 = [
        "hurs",  # surface relative humidity [%]
        "tas",  # near-surface air temp [K]
    ]
    atmo1 = get_struct("atm", atmo_vars1)
    atmo2 = get_struct("atm", atmo_vars2)
    sst = get_struct("sst", ["analysed_sst"])
    [atmo1, atmo2, sst], time_combined = Struct.reinterp_at_finest_grid(
        [atmo1, atmo2, sst], bounds=bounds
    )

    time_combined = clip_inclusive(time_combined, *timelaps)
    csr = atmo1.compute_csr_cg_2d(mesh2d)
    atmo2.compute_csr_cg_2d(mesh2d)
    sst.compute_csr_cg_2d(mesh2d)

    with slim.DataFileWriter(
        prepro_dir + "slim4_atmo", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer_atmo:
        for i, t in enumerate(time_combined):
            [
                (_, data_lat_heat_flux_up),
                (_, data_sens_heat_flux_up),
                (_, shortwave_down),
                (_, shortwave_up),
                (_, longwave_down),
                (_, longwave_up),
            ] = atmo1.get_data(csr, t, ["hfls", "hfss", "rsds", "rsus", "rlds", "rlus"])

            rho = 1000
            shortwave_pen_ratio = 0.76 * 0.302
            lat_heat_flux_wm2 = -data_lat_heat_flux_up  # [W/m2]
            sens_heat_flux_wm2 = -data_sens_heat_flux_up  # [W/m2]
            longwave_heat_fluw_wm2 = longwave_down - longwave_up  # [W/m2]
            shorwave_heat_fluw_wm2 = shortwave_down - shortwave_up

            pen_flux = (
                shorwave_heat_fluw_wm2 * shortwave_pen_ratio / (4180 * rho)
            )  # [K m/s]

            surf_flux = (
                lat_heat_flux_wm2
                + sens_heat_flux_wm2
                + longwave_heat_fluw_wm2
                + shorwave_heat_fluw_wm2 * (1 - shortwave_pen_ratio)
            ) / (4180 * rho)  # [K m/s]
            surf_flux = np.vstack((surf_flux, np.zeros_like(surf_flux)))
            pdist = 19.1  # Fitted to match average absorbance
            pen_dist = np.full_like(pen_flux, pdist)

            # Compute mixing ratios used for latent heat flux corrections
            [(_, data_airtemp_K)] = atmo2.get_data(csr, t, ["tas"])
            [(_, data_relhum_percent)] = atmo2.get_data(csr, t, ["hurs"])
            [(_, data_sst_K)] = sst.get_data(csr, t, ["analysed_sst"])

            airtemp_C = data_airtemp_K - 273.15  # [°C]
            sst_C = data_sst_K - 273.15  # [°C]

            # [(_, data_dewpoint_K)] = atmo2.get_data(csr, t, ["d2m"])
            # dewpoint_C = data_dewpoint_K - 273.15  # [°C]
            # rel_hum = rel_hu_from_dew(dewpoint_C, airtemp_C) # [0-1]
            rel_hum = data_relhum_percent / 100  # [0-1]

            vapor_press_tair = vapor_pressure(airtemp_C)
            vapor_press_sst = vapor_pressure(sst_C)
            air_mix_ratio = (
                rel_hum * 0.622 * vapor_press_tair / (1013 - 0.378 * vapor_press_tair)
            )
            sat_mix_ratio = (
                0.98 * 0.622 * vapor_press_sst / (1013 - 0.378 * vapor_press_sst)
            )

            writer_atmo.append(
                [
                    ("sst", sst_C),
                    ("airtemp", airtemp_C),
                    ("surf_flux", surf_flux),
                    ("pen_flux", pen_flux),
                    ("q_lat_wm2", lat_heat_flux_wm2),
                    ("q_sens_wm2", sens_heat_flux_wm2),
                    ("pen_dist", pen_dist),
                    ("air_mix_ratio", air_mix_ratio),
                    ("sat_mix_ratio", sat_mix_ratio),
                ],
                t,
            )

    mpi.barrier()


def prepro_freshwater_flux():
    """Convert precipitation and evaporation products to freshwater flux."""
    log("Processing Freshwater Flux...")
    atmo_vars1 = [
        "hfls",  # surface upward latent heat flux [W/m²]
        "pr",  # precipitation [kg/m²/s]
    ]
    atmo_vars2 = [
        "ts",  # surface skin temp [K]
    ]
    atmo1 = get_struct("atm", atmo_vars1)
    atmo2 = get_struct("atm", atmo_vars2)
    [atmo1, atmo2], time_combined = Struct.reinterp_at_finest_grid(
        [atmo1, atmo2], bounds=bounds
    )

    time_combined = clip_inclusive(time_combined, *timelaps)

    csr = atmo1.compute_csr_cg_2d(mesh2d)
    atmo2.compute_csr_cg_2d(mesh2d)
    with slim.DataFileWriter(
        prepro_dir + "slim4_fw", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_combined):
            [(_, data_tp), (_, data_lat_heat_flux)] = atmo1.get_data(
                csr, t, ["pr", "hfls"]
            )
            [(_, data_skintemp_K)] = atmo2.get_data(csr, t, ["ts"])
            lat_heat_flux_wm2 = -data_lat_heat_flux  # [W/m2]

            data_skintemp_K = data_skintemp_K
            rho = 1000
            q_lat = lat_heat_flux_wm2 / rho  # [Wm/kg]
            prec = data_tp / rho  # convert from kg/m2/s to m/s
            lat_heat_vap = (
                1.91846e6 * data_skintemp_K**2 / (data_skintemp_K - 33.91) ** 2
            )
            evap = q_lat / lat_heat_vap
            qwater = prec + evap  # [m/s]
            writer.append([("fw", qwater)], t)

    mpi.barrier()


# %%
# We prepare the 3D initial conditions and open-boundary conditions. Currents
# are combined with tidal velocities and rotated to mesh coordinates, while
# temperature and salinity are interpolated from the external ocean product.
def prepro_currents():
    """Prepare 3D velocity initial and open-boundary conditions."""
    log("Processing Currents...")
    uvgrid = get_struct("hydro3d", ["u", "v"])
    tides = tpxo.GridTPXO(
        mesh2d, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone"
    )
    ut = tides.tides["u"]
    vt = tides.tides["v"]

    # Initial condition
    csr_init = uvgrid.compute_csr_cg_3d(mesh, z)
    pnts = uvgrid.get_lonlat_grid(csr_init)
    rot_uv = VectorRotation(mesh2d, pnts)
    ampu, phau = ut.compute_amp_pha(pnts.T)
    ampv, phav = vt.compute_amp_pha(pnts.T)

    path_init = f"{prepro_dir}/slim4_uv_init"
    with slim.DataFileWriter(
        path_init, "CG_INTERP", csr_init.columns, csr_init.row_start, csr_init.weights
    ) as writer_init:
        for i, t in enumerate(timelaps[0:1]):
            [(_, u_base), (_, v_base)] = uvgrid.get_data(csr_init, t, ["u", "v"])
            u_tide = ut.at_points_precomputed(ampu, phau, t)
            v_tide = vt.at_points_precomputed(ampv, phav, t)
            u = u_base + u_tide * 0.01
            v = v_base + v_tide * 0.01
            u, v = rot_uv.rotate(u, v)
            uv = np.vstack([u, v])
            writer_init.append([("u", uv)], t)

    # Boundary conditions
    csr = uvgrid.compute_csr_cg_3d(mesh, z, boundaries=param.open_tags)
    pnts = uvgrid.get_lonlat_grid(csr)
    rot_uv = VectorRotation(mesh2d, pnts)
    ampu, phau = ut.compute_amp_pha(pnts.T)
    ampv, phav = vt.compute_amp_pha(pnts.T)

    with slim.DataFileWriter(
        prepro_dir + "slim4_uv", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_tides):
            [(_, u_base), (_, v_base)] = uvgrid.get_data(csr, t, ["u", "v"])
            u_tide = ut.at_points_precomputed(ampu, phau, t)
            v_tide = vt.at_points_precomputed(ampv, phav, t)
            u = u_base + u_tide * 0.01
            v = v_base + v_tide * 0.01
            u, v = rot_uv.rotate(u, v)
            uv = np.vstack([u, v])
            writer.append([("u", uv)], t)

    mpi.barrier()


def prepro_temp_sal():
    """Prepare temperature and salinity initial and boundary conditions."""
    log("Processing Temperature and Salinity...")
    tsgrid = get_struct("hydro3d", ["temp", "salt"])

    # Initial condition
    csr_init = tsgrid.compute_csr_cg_3d(mesh, z)

    path_init = f"{prepro_dir}/slim4_temp_sal_init"
    with slim.DataFileWriter(
        path_init, "CG_INTERP", csr_init.columns, csr_init.row_start, csr_init.weights
    ) as writer_init:
        for i, t in enumerate(timelaps[0:1]):
            [(_, temp), (_, salt)] = tsgrid.get_data(csr_init, t, ["temp", "salt"])
            writer_init.append([("temp_sal", np.vstack([temp, salt]))], t)

    # Boundary condition
    csr = tsgrid.compute_csr_cg_3d(mesh, z, param.open_tags)

    path_open = f"{prepro_dir}/slim4_temp_sal_open"
    with slim.DataFileWriter(
        path_open, "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(clip_inclusive(tsgrid.time(), *timelaps)):
            [(_, temp), (_, salt)] = tsgrid.get_data(csr, t, ["temp", "salt"])
            writer.append([("temp_sal", np.vstack([temp, salt]))], t)

    mpi.barrier()


# %%
# Finally, we execute all preprocessing steps.

prepro_coriolis()
prepro_reef_denisty()
prepro_wind()
prepro_elevation()
prepro_atmo()
prepro_freshwater_flux()
prepro_currents()
prepro_temp_sal()
log("Preprocessing Done =)")