from veros.core.operators import numpy as npx

from veros import veros_routine, veros_kernel, KernelOutput
from veros.variables import allocate
from veros.core import numerics, utilities, isoneutral
from veros.core.operators import update, update_add, at


@veros_kernel
def explicit_vert_friction(state):
    """
    explicit vertical friction
    dissipation is calculated and added to K_diss_v
    """
    vs = state.variables

    diss = allocate(state.dimensions, ("xt", "yu", "zt"))
    flux_top = allocate(state.dimensions, ("xt", "yu", "zt"))

    """
    vertical friction of zonal momentum
    """
    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[2:-1, 1:-2, :-1])
    flux_top = update(
        flux_top,
        at[1:-2, 1:-2, :-1],
        fxa
        * (vs.u[1:-2, 1:-2, 1:, vs.tau] - vs.u[1:-2, 1:-2, :-1, vs.tau])
        / vs.dzw[npx.newaxis, npx.newaxis, :-1]
        * vs.maskU[1:-2, 1:-2, 1:]
        * vs.maskU[1:-2, 1:-2, :-1],
    )
    flux_top = update(flux_top, at[:, :, -1], 0.0)
    vs.du_mix = update(vs.du_mix, at[:, :, 0], flux_top[:, :, 0] / vs.dzt[0] * vs.maskU[:, :, 0])
    vs.du_mix = update(
        vs.du_mix, at[:, :, 1:], (flux_top[:, :, 1:] - flux_top[:, :, :-1]) / vs.dzt[1:] * vs.maskU[:, :, 1:]
    )

    """
    diagnose dissipation by vertical friction of zonal momentum
    """
    diss = update(
        diss,
        at[1:-2, 1:-2, :-1],
        (vs.u[1:-2, 1:-2, 1:, vs.tau] - vs.u[1:-2, 1:-2, :-1, vs.tau])
        * flux_top[1:-2, 1:-2, :-1]
        / vs.dzw[npx.newaxis, npx.newaxis, :-1],
    )
    diss = update(diss, at[:, :, -1], 0.0)
    diss = numerics.ugrid_to_tgrid(state, diss)
    vs.K_diss_v = vs.K_diss_v + diss

    """
    vertical friction of meridional momentum
    """
    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[1:-2, 2:-1, :-1])
    flux_top = update(
        flux_top,
        at[1:-2, 1:-2, :-1],
        fxa
        * (vs.v[1:-2, 1:-2, 1:, vs.tau] - vs.v[1:-2, 1:-2, :-1, vs.tau])
        / vs.dzw[npx.newaxis, npx.newaxis, :-1]
        * vs.maskV[1:-2, 1:-2, 1:]
        * vs.maskV[1:-2, 1:-2, :-1],
    )
    flux_top = update(flux_top, at[:, :, -1], 0.0)
    vs.dv_mix = update(
        vs.dv_mix,
        at[:, :, 1:],
        (flux_top[:, :, 1:] - flux_top[:, :, :-1]) / vs.dzt[npx.newaxis, npx.newaxis, 1:] * vs.maskV[:, :, 1:],
    )
    vs.dv_mix = update(vs.dv_mix, at[:, :, 0], flux_top[:, :, 0] / vs.dzt[0] * vs.maskV[:, :, 0])

    """
    diagnose dissipation by vertical friction of meridional momentum
    """
    diss = update(
        diss,
        at[1:-2, 1:-2, :-1],
        (vs.v[1:-2, 1:-2, 1:, vs.tau] - vs.v[1:-2, 1:-2, :-1, vs.tau])
        * flux_top[1:-2, 1:-2, :-1]
        / vs.dzw[npx.newaxis, npx.newaxis, :-1],
    )
    diss = update(diss, at[:, :, -1], 0.0)
    diss = numerics.vgrid_to_tgrid(state, diss)
    vs.K_diss_v = vs.K_diss_v + diss

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_v=vs.K_diss_v)


@veros_kernel
def implicit_vert_friction(state):
    """
    vertical friction
    dissipation is calculated and added to K_diss_v
    """
    vs = state.variables
    settings = state.settings

    diss = allocate(state.dimensions, ("xt", "yu", "zt"))
    a_tri = allocate(state.dimensions, ("xt", "yu", "zt"))[1:-2, 1:-2]
    b_tri = allocate(state.dimensions, ("xt", "yu", "zt"))[1:-2, 1:-2]
    c_tri = allocate(state.dimensions, ("xt", "yu", "zt"))[1:-2, 1:-2]
    d_tri = allocate(state.dimensions, ("xt", "yu", "zt"))[1:-2, 1:-2]
    delta = allocate(state.dimensions, ("xt", "yu", "zt"))[1:-2, 1:-2]
    flux_top = allocate(state.dimensions, ("xt", "yu", "zt"))

    """
    implicit vertical friction of zonal momentum
    """
    kss = npx.maximum(vs.kbot[1:-2, 1:-2], vs.kbot[2:-1, 1:-2])
    _, water_mask, edge_mask = utilities.create_water_masks(kss, settings.nz)

    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[2:-1, 1:-2, :-1])
    delta = update(
        delta, at[:, :, :-1], settings.dt_mom / vs.dzw[:-1] * fxa * vs.maskU[1:-2, 1:-2, 1:] * vs.maskU[1:-2, 1:-2, :-1]
    )
    a_tri = update(a_tri, at[:, :, 1:], -delta[:, :, :-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:])
    b_tri = update(b_tri, at[:, :, 1:], 1 + delta[:, :, :-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:])
    b_tri = update_add(b_tri, at[:, :, 1:-1], delta[:, :, 1:-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:-1])
    b_tri_edge = 1 + delta / vs.dzt[npx.newaxis, npx.newaxis, :]
    c_tri = update(c_tri, at[...], -delta / vs.dzt[npx.newaxis, npx.newaxis, :])
    d_tri = update(d_tri, at[...], vs.u[1:-2, 1:-2, :, vs.tau])

    res = utilities.solve_implicit(a_tri, b_tri, c_tri, d_tri, water_mask, b_edge=b_tri_edge, edge_mask=edge_mask)
    vs.u = update(vs.u, at[1:-2, 1:-2, :, vs.taup1], npx.where(water_mask, res, vs.u[1:-2, 1:-2, :, vs.taup1]))
    vs.du_mix = update(
        vs.du_mix, at[1:-2, 1:-2], (vs.u[1:-2, 1:-2, :, vs.taup1] - vs.u[1:-2, 1:-2, :, vs.tau]) / settings.dt_mom
    )

    """
    diagnose dissipation by vertical friction of zonal momentum
    """
    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[2:-1, 1:-2, :-1])
    flux_top = update(
        flux_top,
        at[1:-2, 1:-2, :-1],
        fxa
        * (vs.u[1:-2, 1:-2, 1:, vs.taup1] - vs.u[1:-2, 1:-2, :-1, vs.taup1])
        / vs.dzw[:-1]
        * vs.maskU[1:-2, 1:-2, 1:]
        * vs.maskU[1:-2, 1:-2, :-1],
    )
    diss = update(
        diss,
        at[1:-2, 1:-2, :-1],
        (vs.u[1:-2, 1:-2, 1:, vs.tau] - vs.u[1:-2, 1:-2, :-1, vs.tau]) * flux_top[1:-2, 1:-2, :-1] / vs.dzw[:-1],
    )
    diss = update(diss, at[:, :, -1], 0.0)
    diss = numerics.ugrid_to_tgrid(state, diss)
    vs.K_diss_v = vs.K_diss_v + diss

    """
    implicit vertical friction of meridional momentum
    """
    kss = npx.maximum(vs.kbot[1:-2, 1:-2], vs.kbot[1:-2, 2:-1])
    _, water_mask, edge_mask = utilities.create_water_masks(kss, settings.nz)

    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[1:-2, 2:-1, :-1])
    delta = update(
        delta,
        at[:, :, :-1],
        settings.dt_mom
        / vs.dzw[npx.newaxis, npx.newaxis, :-1]
        * fxa
        * vs.maskV[1:-2, 1:-2, 1:]
        * vs.maskV[1:-2, 1:-2, :-1],
    )
    a_tri = update(a_tri, at[:, :, 1:], -delta[:, :, :-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:])
    b_tri = update(b_tri, at[:, :, 1:], 1 + delta[:, :, :-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:])
    b_tri = update_add(b_tri, at[:, :, 1:-1], delta[:, :, 1:-1] / vs.dzt[npx.newaxis, npx.newaxis, 1:-1])
    b_tri_edge = 1 + delta / vs.dzt[npx.newaxis, npx.newaxis, :]
    c_tri = update(c_tri, at[:, :, :-1], -delta[:, :, :-1] / vs.dzt[npx.newaxis, npx.newaxis, :-1])
    c_tri = update(c_tri, at[:, :, -1], 0.0)
    d_tri = update(d_tri, at[...], vs.v[1:-2, 1:-2, :, vs.tau])

    res = utilities.solve_implicit(a_tri, b_tri, c_tri, d_tri, water_mask, b_edge=b_tri_edge, edge_mask=edge_mask)
    vs.v = update(vs.v, at[1:-2, 1:-2, :, vs.taup1], npx.where(water_mask, res, vs.v[1:-2, 1:-2, :, vs.taup1]))
    vs.dv_mix = update(
        vs.dv_mix, at[1:-2, 1:-2], (vs.v[1:-2, 1:-2, :, vs.taup1] - vs.v[1:-2, 1:-2, :, vs.tau]) / settings.dt_mom
    )

    """
    diagnose dissipation by vertical friction of meridional momentum
    """
    fxa = 0.5 * (vs.kappaM[1:-2, 1:-2, :-1] + vs.kappaM[1:-2, 2:-1, :-1])
    flux_top = update(
        flux_top,
        at[1:-2, 1:-2, :-1],
        fxa
        * (vs.v[1:-2, 1:-2, 1:, vs.taup1] - vs.v[1:-2, 1:-2, :-1, vs.taup1])
        / vs.dzw[:-1]
        * vs.maskV[1:-2, 1:-2, 1:]
        * vs.maskV[1:-2, 1:-2, :-1],
    )
    diss = update(
        diss,
        at[1:-2, 1:-2, :-1],
        (vs.v[1:-2, 1:-2, 1:, vs.tau] - vs.v[1:-2, 1:-2, :-1, vs.tau]) * flux_top[1:-2, 1:-2, :-1] / vs.dzw[:-1],
    )
    diss = update(diss, at[:, :, -1], 0.0)
    diss = numerics.vgrid_to_tgrid(state, diss)
    vs.K_diss_v = vs.K_diss_v + diss

    return KernelOutput(u=vs.u, v=vs.v, du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_v=vs.K_diss_v)


@veros_kernel
def rayleigh_friction(state):
    """
    interior Rayleigh friction
    dissipation is calculated and added to K_diss_bot
    """
    vs = state.variables
    settings = state.settings

    vs.du_mix = update_add(vs.du_mix, at[...], -1 * vs.maskU * settings.r_ray * vs.u[..., vs.tau])
    if settings.enable_conserve_energy:
        diss = vs.maskU * settings.r_ray * vs.u[..., vs.tau] ** 2
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_u(state, diss))

    vs.dv_mix = update_add(vs.dv_mix, at[...], -1 * vs.maskV * settings.r_ray * vs.v[..., vs.tau])
    if settings.enable_conserve_energy:
        diss = vs.maskV * settings.r_ray * vs.v[..., vs.tau] ** 2
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_bot=vs.K_diss_bot)


@veros_kernel
def linear_bottom_friction(state):
    """
    linear bottom friction
    dissipation is calculated and added to K_diss_bot
    """
    vs = state.variables
    settings = state.settings

    if settings.enable_bottom_friction_var:
        """
        with spatially varying coefficient
        """
        k = npx.maximum(vs.kbot[1:-2, 2:-2], vs.kbot[2:-1, 2:-2]) - 1
        mask = npx.arange(settings.nz) == k[:, :, npx.newaxis]
        vs.du_mix = update_add(
            vs.du_mix,
            at[1:-2, 2:-2],
            -(vs.maskU[1:-2, 2:-2] * vs.r_bot_var_u[1:-2, 2:-2, npx.newaxis]) * vs.u[1:-2, 2:-2, :, vs.tau] * mask,
        )
        if settings.enable_conserve_energy:
            diss = allocate(state.dimensions, ("xt", "yu", "zt"))
            diss = update(
                diss,
                at[1:-2, 2:-2],
                vs.maskU[1:-2, 2:-2]
                * vs.r_bot_var_u[1:-2, 2:-2, npx.newaxis]
                * vs.u[1:-2, 2:-2, :, vs.tau] ** 2
                * mask,
            )
            vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_u(state, diss))

        k = npx.maximum(vs.kbot[2:-2, 2:-1], vs.kbot[2:-2, 1:-2]) - 1
        mask = npx.arange(settings.nz) == k[:, :, npx.newaxis]
        vs.dv_mix = update_add(
            vs.dv_mix,
            at[2:-2, 1:-2],
            -(vs.maskV[2:-2, 1:-2] * vs.r_bot_var_v[2:-2, 1:-2, npx.newaxis]) * vs.v[2:-2, 1:-2, :, vs.tau] * mask,
        )
        if settings.enable_conserve_energy:
            diss = allocate(state.dimensions, ("xt", "yu", "zt"))
            diss = update(
                diss,
                at[2:-2, 1:-2],
                vs.maskV[2:-2, 1:-2]
                * vs.r_bot_var_v[2:-2, 1:-2, npx.newaxis]
                * vs.v[2:-2, 1:-2, :, vs.tau] ** 2
                * mask,
            )
            vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_v(state, diss))
    else:
        """
        with constant coefficient
        """
        k = npx.maximum(vs.kbot[1:-2, 2:-2], vs.kbot[2:-1, 2:-2]) - 1
        mask = npx.arange(settings.nz) == k[:, :, npx.newaxis]

        vs.du_mix = update_add(
            vs.du_mix, at[1:-2, 2:-2], -1 * vs.maskU[1:-2, 2:-2] * settings.r_bot * vs.u[1:-2, 2:-2, :, vs.tau] * mask
        )
        if settings.enable_conserve_energy:
            diss = allocate(state.dimensions, ("xt", "yu", "zt"))
            diss = update(
                diss, at[1:-2, 2:-2], vs.maskU[1:-2, 2:-2] * settings.r_bot * vs.u[1:-2, 2:-2, :, vs.tau] ** 2 * mask
            )
            vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_u(state, diss))

        k = npx.maximum(vs.kbot[2:-2, 2:-1], vs.kbot[2:-2, 1:-2]) - 1
        mask = npx.arange(settings.nz) == k[:, :, npx.newaxis]

        vs.dv_mix = update_add(
            vs.dv_mix, at[2:-2, 1:-2], -1 * vs.maskV[2:-2, 1:-2] * settings.r_bot * vs.v[2:-2, 1:-2, :, vs.tau] * mask
        )
        if settings.enable_conserve_energy:
            diss = allocate(state.dimensions, ("xt", "yu", "zt"))
            diss = update(
                diss, at[2:-2, 1:-2], vs.maskV[2:-2, 1:-2] * settings.r_bot * vs.v[2:-2, 1:-2, :, vs.tau] ** 2 * mask
            )
            vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_bot=vs.K_diss_bot)


@veros_kernel
def quadratic_bottom_friction(state):
    """
    quadratic bottom friction
    dissipation is calculated and added to K_diss_bot
    """
    vs = state.variables
    settings = state.settings

    # we might want to account for EKE in the drag, also a tidal residual
    k = npx.maximum(vs.kbot[1:-2, 2:-2], vs.kbot[2:-1, 2:-2]) - 1
    mask = k[..., npx.newaxis] == npx.arange(settings.nz)[npx.newaxis, npx.newaxis, :]
    fxa = (
        vs.maskV[1:-2, 2:-2, :] * vs.v[1:-2, 2:-2, :, vs.tau] ** 2
        + vs.maskV[1:-2, 1:-3, :] * vs.v[1:-2, 1:-3, :, vs.tau] ** 2
        + vs.maskV[2:-1, 2:-2, :] * vs.v[2:-1, 2:-2, :, vs.tau] ** 2
        + vs.maskV[2:-1, 1:-3, :] * vs.v[2:-1, 1:-3, :, vs.tau] ** 2
    )
    fxa = npx.sqrt(vs.u[1:-2, 2:-2, :, vs.tau] ** 2 + 0.25 * fxa)
    aloc = (
        vs.maskU[1:-2, 2:-2, :]
        * settings.r_quad_bot
        * vs.u[1:-2, 2:-2, :, vs.tau]
        * fxa
        / vs.dzt[npx.newaxis, npx.newaxis, :]
        * mask
    )
    vs.du_mix = update_add(vs.du_mix, at[1:-2, 2:-2, :], -aloc)

    if settings.enable_conserve_energy:
        diss = allocate(state.dimensions, ("xt", "yu", "zt"))
        diss = update(diss, at[1:-2, 2:-2, :], aloc * vs.u[1:-2, 2:-2, :, vs.tau])
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_u(state, diss))

    k = npx.maximum(vs.kbot[2:-2, 1:-2], vs.kbot[2:-2, 2:-1]) - 1
    mask = k[..., npx.newaxis] == npx.arange(settings.nz)[npx.newaxis, npx.newaxis, :]
    fxa = (
        vs.maskU[2:-2, 1:-2, :] * vs.u[2:-2, 1:-2, :, vs.tau] ** 2
        + vs.maskU[1:-3, 1:-2, :] * vs.u[1:-3, 1:-2, :, vs.tau] ** 2
        + vs.maskU[2:-2, 2:-1, :] * vs.u[2:-2, 2:-1, :, vs.tau] ** 2
        + vs.maskU[1:-3, 2:-1, :] * vs.u[1:-3, 2:-1, :, vs.tau] ** 2
    )
    fxa = npx.sqrt(vs.v[2:-2, 1:-2, :, vs.tau] ** 2 + 0.25 * fxa)
    aloc = (
        vs.maskV[2:-2, 1:-2, :]
        * settings.r_quad_bot
        * vs.v[2:-2, 1:-2, :, vs.tau]
        * fxa
        / vs.dzt[npx.newaxis, npx.newaxis, :]
        * mask
    )
    vs.dv_mix = update_add(vs.dv_mix, at[2:-2, 1:-2, :], -aloc)

    if settings.enable_conserve_energy:
        diss = allocate(state.dimensions, ("xt", "yu", "zt"))
        diss = update(diss, at[2:-2, 1:-2, :], aloc * vs.v[2:-2, 1:-2, :, vs.tau])
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_bot=vs.K_diss_bot)


@veros_kernel
def harmonic_friction(state):
    """
    horizontal harmonic friction
    dissipation is calculated and added to K_diss_h
    """
    vs = state.variables
    settings = state.settings

    diss = allocate(state.dimensions, ("xt", "yt", "zt"))
    flux_east = allocate(state.dimensions, ("xu", "yt", "zt"))
    flux_north = allocate(state.dimensions, ("xt", "yu", "zt"))

    """
    Zonal velocity
    """
    if settings.enable_hor_friction_cos_scaling:
        fxa = vs.cost**settings.hor_friction_cosPower
        flux_east = update(
            flux_east,
            at[:-1],
            settings.A_h
            * fxa[npx.newaxis, :, npx.newaxis]
            * (vs.u[1:, :, :, vs.tau] - vs.u[:-1, :, :, vs.tau])
            / (vs.cost * vs.dxt[1:, npx.newaxis])[:, :, npx.newaxis]
            * vs.maskU[1:]
            * vs.maskU[:-1],
        )
        fxa = vs.cosu**settings.hor_friction_cosPower
        flux_north = update(
            flux_north,
            at[:, :-1],
            settings.A_h
            * fxa[npx.newaxis, :-1, npx.newaxis]
            * (vs.u[:, 1:, :, vs.tau] - vs.u[:, :-1, :, vs.tau])
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * vs.maskU[:, 1:]
            * vs.maskU[:, :-1]
            * vs.cosu[npx.newaxis, :-1, npx.newaxis],
        )
        if settings.enable_noslip_lateral:
            flux_north = update_add(
                flux_north,
                at[:, :-1],
                2
                * settings.A_h
                * fxa[npx.newaxis, :-1, npx.newaxis]
                * (vs.u[:, 1:, :, vs.tau])
                / vs.dyu[npx.newaxis, :-1, npx.newaxis]
                * vs.maskU[:, 1:]
                * (1 - vs.maskU[:, :-1])
                * vs.cosu[npx.newaxis, :-1, npx.newaxis]
                - 2
                * settings.A_h
                * fxa[npx.newaxis, :-1, npx.newaxis]
                * (vs.u[:, :-1, :, vs.tau])
                / vs.dyu[npx.newaxis, :-1, npx.newaxis]
                * (1 - vs.maskU[:, 1:])
                * vs.maskU[:, :-1]
                * vs.cosu[npx.newaxis, :-1, npx.newaxis],
            )
    else:
        flux_east = update(
            flux_east,
            at[:-1, :, :],
            settings.A_h
            * (vs.u[1:, :, :, vs.tau] - vs.u[:-1, :, :, vs.tau])
            / (vs.cost * vs.dxt[1:, npx.newaxis])[:, :, npx.newaxis]
            * vs.maskU[1:]
            * vs.maskU[:-1],
        )
        flux_north = update(
            flux_north,
            at[:, :-1, :],
            settings.A_h
            * (vs.u[:, 1:, :, vs.tau] - vs.u[:, :-1, :, vs.tau])
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * vs.maskU[:, 1:]
            * vs.maskU[:, :-1]
            * vs.cosu[npx.newaxis, :-1, npx.newaxis],
        )
        if settings.enable_noslip_lateral:
            flux_north = update_add(
                flux_north,
                at[:, :-1],
                2
                * settings.A_h
                * vs.u[:, 1:, :, vs.tau]
                / vs.dyu[npx.newaxis, :-1, npx.newaxis]
                * vs.maskU[:, 1:]
                * (1 - vs.maskU[:, :-1])
                * vs.cosu[npx.newaxis, :-1, npx.newaxis]
                - 2
                * settings.A_h
                * vs.u[:, :-1, :, vs.tau]
                / vs.dyu[npx.newaxis, :-1, npx.newaxis]
                * (1 - vs.maskU[:, 1:])
                * vs.maskU[:, :-1]
                * vs.cosu[npx.newaxis, :-1, npx.newaxis],
            )

    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    """
    update tendency
    """
    vs.du_mix = update_add(
        vs.du_mix,
        at[2:-2, 2:-2, :],
        vs.maskU[2:-2, 2:-2]
        * (
            (flux_east[2:-2, 2:-2] - flux_east[1:-3, 2:-2])
            / (vs.cost[2:-2] * vs.dxu[2:-2, npx.newaxis])[:, :, npx.newaxis]
            + (flux_north[2:-2, 2:-2] - flux_north[2:-2, 1:-3])
            / (vs.cost[2:-2] * vs.dyt[2:-2])[npx.newaxis, :, npx.newaxis]
        ),
    )

    if settings.enable_conserve_energy:
        """
        diagnose dissipation by lateral friction
        """
        diss = update(
            diss,
            at[1:-2, 2:-2],
            0.5
            * (
                (vs.u[2:-1, 2:-2, :, vs.tau] - vs.u[1:-2, 2:-2, :, vs.tau]) * flux_east[1:-2, 2:-2]
                + (vs.u[1:-2, 2:-2, :, vs.tau] - vs.u[:-3, 2:-2, :, vs.tau]) * flux_east[:-3, 2:-2]
            )
            / (vs.cost[2:-2] * vs.dxu[1:-2, npx.newaxis])[:, :, npx.newaxis]
            + 0.5
            * (
                (vs.u[1:-2, 3:-1, :, vs.tau] - vs.u[1:-2, 2:-2, :, vs.tau]) * flux_north[1:-2, 2:-2]
                + (vs.u[1:-2, 2:-2, :, vs.tau] - vs.u[1:-2, 1:-3, :, vs.tau]) * flux_north[1:-2, 1:-3]
            )
            / (vs.cost[2:-2] * vs.dyt[2:-2])[npx.newaxis, :, npx.newaxis],
        )
        vs.K_diss_h = numerics.calc_diss_u(state, diss)

    """
    Meridional velocity
    """
    if settings.enable_hor_friction_cos_scaling:
        flux_east = update(
            flux_east,
            at[:-1],
            settings.A_h
            * vs.cosu[npx.newaxis, :, npx.newaxis] ** settings.hor_friction_cosPower
            * (vs.v[1:, :, :, vs.tau] - vs.v[:-1, :, :, vs.tau])
            / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
            * vs.maskV[1:]
            * vs.maskV[:-1],
        )

        if settings.enable_noslip_lateral:
            flux_east = update_add(
                flux_east,
                at[:-1],
                2
                * settings.A_h
                * fxa[npx.newaxis, :, npx.newaxis]
                * vs.v[1:, :, :, vs.tau]
                / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
                * vs.maskV[1:]
                * (1 - vs.maskV[:-1])
                - 2
                * settings.A_h
                * fxa[npx.newaxis, :, npx.newaxis]
                * vs.v[:-1, :, :, vs.tau]
                / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
                * (1 - vs.maskV[1:])
                * vs.maskV[:-1],
            )

        flux_north = update(
            flux_north,
            at[:, :-1],
            settings.A_h
            * vs.cost[npx.newaxis, 1:, npx.newaxis] ** settings.hor_friction_cosPower
            * (vs.v[:, 1:, :, vs.tau] - vs.v[:, :-1, :, vs.tau])
            / vs.dyt[npx.newaxis, 1:, npx.newaxis]
            * vs.cost[npx.newaxis, 1:, npx.newaxis]
            * vs.maskV[:, :-1]
            * vs.maskV[:, 1:],
        )
    else:
        flux_east = update(
            flux_east,
            at[:-1],
            settings.A_h
            * (vs.v[1:, :, :, vs.tau] - vs.v[:-1, :, :, vs.tau])
            / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
            * vs.maskV[1:]
            * vs.maskV[:-1],
        )

        if settings.enable_noslip_lateral:
            flux_east = update_add(
                flux_east,
                at[:-1],
                2
                * settings.A_h
                * vs.v[1:, :, :, vs.tau]
                / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
                * vs.maskV[1:]
                * (1 - vs.maskV[:-1])
                - 2
                * settings.A_h
                * vs.v[:-1, :, :, vs.tau]
                / (vs.cosu * vs.dxu[:-1, npx.newaxis])[:, :, npx.newaxis]
                * (1 - vs.maskV[1:])
                * vs.maskV[:-1],
            )

        flux_north = update(
            flux_north,
            at[:, :-1],
            settings.A_h
            * (vs.v[:, 1:, :, vs.tau] - vs.v[:, :-1, :, vs.tau])
            / vs.dyt[npx.newaxis, 1:, npx.newaxis]
            * vs.cost[npx.newaxis, 1:, npx.newaxis]
            * vs.maskV[:, :-1]
            * vs.maskV[:, 1:],
        )

    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    """
    update tendency
    """
    vs.dv_mix = update_add(
        vs.dv_mix,
        at[2:-2, 2:-2],
        vs.maskV[2:-2, 2:-2]
        * (
            (flux_east[2:-2, 2:-2] - flux_east[1:-3, 2:-2])
            / (vs.cosu[2:-2] * vs.dxt[2:-2, npx.newaxis])[:, :, npx.newaxis]
            + (flux_north[2:-2, 2:-2] - flux_north[2:-2, 1:-3])
            / (vs.dyu[2:-2] * vs.cosu[2:-2])[npx.newaxis, :, npx.newaxis]
        ),
    )

    if settings.enable_conserve_energy:
        """
        diagnose dissipation by lateral friction
        """
        diss = update(
            diss,
            at[2:-2, 1:-2],
            0.5
            * (
                (vs.v[3:-1, 1:-2, :, vs.tau] - vs.v[2:-2, 1:-2, :, vs.tau]) * flux_east[2:-2, 1:-2]
                + (vs.v[2:-2, 1:-2, :, vs.tau] - vs.v[1:-3, 1:-2, :, vs.tau]) * flux_east[1:-3, 1:-2]
            )
            / (vs.cosu[1:-2] * vs.dxt[2:-2, npx.newaxis])[:, :, npx.newaxis]
            + 0.5
            * (
                (vs.v[2:-2, 2:-1, :, vs.tau] - vs.v[2:-2, 1:-2, :, vs.tau]) * flux_north[2:-2, 1:-2]
                + (vs.v[2:-2, 1:-2, :, vs.tau] - vs.v[2:-2, :-3, :, vs.tau]) * flux_north[2:-2, :-3]
            )
            / (vs.cosu[1:-2] * vs.dyu[1:-2])[npx.newaxis, :, npx.newaxis],
        )
        vs.K_diss_h = update_add(vs.K_diss_h, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_h=vs.K_diss_h)


@veros_kernel
def biharmonic_friction(state):
    """
    horizontal biharmonic friction
    dissipation is calculated and added to K_diss_h
    """
    vs = state.variables
    settings = state.settings

    flux_east = allocate(state.dimensions, ("xu", "yt", "zt"))
    flux_north = allocate(state.dimensions, ("xt", "yu", "zt"))
    visc = npx.sqrt(abs(settings.A_hbi))

    # each of these enters twice, so we halve the power
    cost_scaled = vs.cost ** (0.5 * settings.biharmonic_friction_cosPower)
    cosu_scaled = vs.cosu ** (0.5 * settings.biharmonic_friction_cosPower)

    """
    Zonal velocity
    """
    flux_east = update(
        flux_east,
        at[:-1, :, :],
        visc
        * cost_scaled[npx.newaxis, :, npx.newaxis]
        * (vs.u[1:, :, :, vs.tau] - vs.u[:-1, :, :, vs.tau])
        / (vs.cost[npx.newaxis, :, npx.newaxis] * vs.dxt[1:, npx.newaxis, npx.newaxis])
        * vs.maskU[1:, :, :]
        * vs.maskU[:-1, :, :],
    )
    flux_north = update(
        flux_north,
        at[:, :-1, :],
        visc
        * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
        * (vs.u[:, 1:, :, vs.tau] - vs.u[:, :-1, :, vs.tau])
        / vs.dyu[npx.newaxis, :-1, npx.newaxis]
        * vs.maskU[:, 1:, :]
        * vs.maskU[:, :-1, :]
        * vs.cosu[npx.newaxis, :-1, npx.newaxis],
    )

    if settings.enable_noslip_lateral:
        flux_north = update_add(
            flux_north,
            at[:, :-1],
            2
            * visc
            * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
            * vs.u[:, 1:, :, vs.tau]
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * vs.maskU[:, 1:]
            * (1 - vs.maskU[:, :-1])
            * vs.cosu[npx.newaxis, :-1, npx.newaxis]
            - 2
            * visc
            * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
            * vs.u[:, :-1, :, vs.tau]
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * (1 - vs.maskU[:, 1:])
            * vs.maskU[:, :-1]
            * vs.cosu[npx.newaxis, :-1, npx.newaxis],
        )

    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    del2 = allocate(state.dimensions, ("xt", "yu", "zt"))
    del2 = update(
        del2,
        at[1:, 1:, :],
        (flux_east[1:, 1:, :] - flux_east[:-1, 1:, :])
        / (vs.cost[npx.newaxis, 1:, npx.newaxis] * vs.dxu[1:, npx.newaxis, npx.newaxis])
        + (flux_north[1:, 1:, :] - flux_north[1:, :-1, :])
        / (vs.cost[npx.newaxis, 1:, npx.newaxis] * vs.dyt[npx.newaxis, 1:, npx.newaxis]),
    )

    flux_east = update(
        flux_east,
        at[:-1, :, :],
        visc
        * cost_scaled[npx.newaxis, :, npx.newaxis]
        * (del2[1:, :, :] - del2[:-1, :, :])
        / (vs.cost[npx.newaxis, :, npx.newaxis] * vs.dxt[1:, npx.newaxis, npx.newaxis])
        * vs.maskU[1:, :, :]
        * vs.maskU[:-1, :, :],
    )
    flux_north = update(
        flux_north,
        at[:, :-1, :],
        visc
        * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
        * (del2[:, 1:, :] - del2[:, :-1, :])
        / vs.dyu[npx.newaxis, :-1, npx.newaxis]
        * vs.maskU[:, 1:, :]
        * vs.maskU[:, :-1, :]
        * vs.cosu[npx.newaxis, :-1, npx.newaxis],
    )

    if settings.enable_noslip_lateral:
        flux_north = update_add(
            flux_north,
            at[:, :-1, :],
            2
            * visc
            * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
            * del2[:, 1:, :]
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * vs.maskU[:, 1:, :]
            * (1 - vs.maskU[:, :-1, :])
            * vs.cosu[npx.newaxis, :-1, npx.newaxis]
            - 2
            * visc
            * cosu_scaled[npx.newaxis, :-1, npx.newaxis]
            * del2[:, :-1, :]
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * (1 - vs.maskU[:, 1:, :])
            * vs.maskU[:, :-1, :]
            * vs.cosu[npx.newaxis, :-1, npx.newaxis],
        )

    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    """
    update tendency
    """
    vs.du_mix = update_add(
        vs.du_mix,
        at[2:-2, 2:-2, :],
        -1
        * vs.maskU[2:-2, 2:-2, :]
        * (
            (flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
            / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dxu[2:-2, npx.newaxis, npx.newaxis])
            + (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
            / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dyt[npx.newaxis, 2:-2, npx.newaxis])
        ),
    )

    if settings.enable_conserve_energy:
        """
        diagnose dissipation by lateral friction
        """
        flux_east = utilities.enforce_boundaries(flux_east, settings.enable_cyclic_x)
        flux_north = utilities.enforce_boundaries(flux_north, settings.enable_cyclic_x)
        diss = allocate(state.dimensions, ("xt", "yu", "zt"))
        diss = update(
            diss,
            at[1:-2, 2:-2, :],
            -0.5
            * (
                (vs.u[2:-1, 2:-2, :, vs.tau] - vs.u[1:-2, 2:-2, :, vs.tau]) * flux_east[1:-2, 2:-2, :]
                + (vs.u[1:-2, 2:-2, :, vs.tau] - vs.u[:-3, 2:-2, :, vs.tau]) * flux_east[:-3, 2:-2, :]
            )
            / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dxu[1:-2, npx.newaxis, npx.newaxis])
            - 0.5
            * (
                (vs.u[1:-2, 3:-1, :, vs.tau] - vs.u[1:-2, 2:-2, :, vs.tau]) * flux_north[1:-2, 2:-2, :]
                + (vs.u[1:-2, 2:-2, :, vs.tau] - vs.u[1:-2, 1:-3, :, vs.tau]) * flux_north[1:-2, 1:-3, :]
            )
            / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dyt[npx.newaxis, 2:-2, npx.newaxis]),
        )
        vs.K_diss_h = numerics.calc_diss_u(state, diss)

    """
    Meridional velocity
    """
    flux_east = update(
        flux_east,
        at[:-1, :, :],
        visc
        * cosu_scaled[npx.newaxis, :, npx.newaxis]
        * (vs.v[1:, :, :, vs.tau] - vs.v[:-1, :, :, vs.tau])
        / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
        * vs.maskV[1:, :, :]
        * vs.maskV[:-1, :, :],
    )

    if settings.enable_noslip_lateral:
        flux_east = update_add(
            flux_east,
            at[:-1, :, :],
            2
            * visc
            * cosu_scaled[npx.newaxis, :, npx.newaxis]
            * vs.v[1:, :, :, vs.tau]
            / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
            * vs.maskV[1:, :, :]
            * (1 - vs.maskV[:-1, :, :])
            - 2
            * visc
            * cosu_scaled[npx.newaxis, :, npx.newaxis]
            * vs.v[:-1, :, :, vs.tau]
            / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
            * (1 - vs.maskV[1:, :, :])
            * vs.maskV[:-1, :, :],
        )

    flux_north = update(
        flux_north,
        at[:, :-1, :],
        visc
        * cost_scaled[npx.newaxis, :-1, npx.newaxis]
        * (vs.v[:, 1:, :, vs.tau] - vs.v[:, :-1, :, vs.tau])
        / vs.dyt[npx.newaxis, 1:, npx.newaxis]
        * vs.cost[npx.newaxis, 1:, npx.newaxis]
        * vs.maskV[:, :-1, :]
        * vs.maskV[:, 1:, :],
    )
    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    del2 = update(
        del2,
        at[1:, 1:, :],
        (flux_east[1:, 1:, :] - flux_east[:-1, 1:, :])
        / (vs.cosu[npx.newaxis, 1:, npx.newaxis] * vs.dxt[1:, npx.newaxis, npx.newaxis])
        + (flux_north[1:, 1:, :] - flux_north[1:, :-1, :])
        / (vs.dyu[npx.newaxis, 1:, npx.newaxis] * vs.cosu[npx.newaxis, 1:, npx.newaxis]),
    )

    flux_east = update(
        flux_east,
        at[:-1, :, :],
        visc
        * cosu_scaled[npx.newaxis, :, npx.newaxis]
        * (del2[1:, :, :] - del2[:-1, :, :])
        / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
        * vs.maskV[1:, :, :]
        * vs.maskV[:-1, :, :],
    )

    if settings.enable_noslip_lateral:
        flux_east = update_add(
            flux_east,
            at[:-1, :, :],
            2
            * visc
            * cosu_scaled[npx.newaxis, :, npx.newaxis]
            * del2[1:, :, :]
            / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
            * vs.maskV[1:, :, :]
            * (1 - vs.maskV[:-1, :, :])
            - 2
            * visc
            * cosu_scaled[npx.newaxis, :, npx.newaxis]
            * del2[:-1, :, :]
            / (vs.cosu[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
            * (1 - vs.maskV[1:, :, :])
            * vs.maskV[:-1, :, :],
        )

    flux_north = update(
        flux_north,
        at[:, :-1, :],
        visc
        * cost_scaled[npx.newaxis, :-1, npx.newaxis]
        * (del2[:, 1:, :] - del2[:, :-1, :])
        / vs.dyt[npx.newaxis, 1:, npx.newaxis]
        * vs.cost[npx.newaxis, 1:, npx.newaxis]
        * vs.maskV[:, :-1, :]
        * vs.maskV[:, 1:, :],
    )
    flux_east = update(flux_east, at[-1, :, :], 0.0)
    flux_north = update(flux_north, at[:, -1, :], 0.0)

    """
    update tendency
    """
    vs.dv_mix = update_add(
        vs.dv_mix,
        at[2:-2, 2:-2, :],
        -1
        * vs.maskV[2:-2, 2:-2, :]
        * (
            (flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
            / (vs.cosu[npx.newaxis, 2:-2, npx.newaxis] * vs.dxt[2:-2, npx.newaxis, npx.newaxis])
            + (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
            / (vs.dyu[npx.newaxis, 2:-2, npx.newaxis] * vs.cosu[npx.newaxis, 2:-2, npx.newaxis])
        ),
    )

    if settings.enable_conserve_energy:
        """
        diagnose dissipation by lateral friction
        """
        flux_east = utilities.enforce_boundaries(flux_east, settings.enable_cyclic_x)
        flux_north = utilities.enforce_boundaries(flux_north, settings.enable_cyclic_x)
        diss = update(
            diss,
            at[2:-2, 1:-2, :],
            -0.5
            * (
                (vs.v[3:-1, 1:-2, :, vs.tau] - vs.v[2:-2, 1:-2, :, vs.tau]) * flux_east[2:-2, 1:-2, :]
                + (vs.v[2:-2, 1:-2, :, vs.tau] - vs.v[1:-3, 1:-2, :, vs.tau]) * flux_east[1:-3, 1:-2, :]
            )
            / (vs.cosu[npx.newaxis, 1:-2, npx.newaxis] * vs.dxt[2:-2, npx.newaxis, npx.newaxis])
            - 0.5
            * (
                (vs.v[2:-2, 2:-1, :, vs.tau] - vs.v[2:-2, 1:-2, :, vs.tau]) * flux_north[2:-2, 1:-2, :]
                + (vs.v[2:-2, 1:-2, :, vs.tau] - vs.v[2:-2, :-3, :, vs.tau]) * flux_north[2:-2, :-3, :]
            )
            / (vs.cosu[npx.newaxis, 1:-2, npx.newaxis] * vs.dyu[npx.newaxis, 1:-2, npx.newaxis]),
        )
        vs.K_diss_h = update_add(vs.K_diss_h, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_h=vs.K_diss_h)


@veros_kernel
def momentum_sources(state):
    """
    other momentum sources
    dissipation is calculated and added to K_diss_bot
    """
    vs = state.variables
    settings = state.settings

    vs.du_mix = update_add(vs.du_mix, at[...], vs.maskU * vs.u_source)
    if settings.enable_conserve_energy:
        diss = -1 * vs.maskU * vs.u[..., vs.tau] * vs.u_source
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_u(state, diss))

    vs.dv_mix = update_add(vs.dv_mix, at[...], vs.maskV * vs.v_source)
    if settings.enable_conserve_energy:
        diss = -1 * vs.maskV * vs.v[..., vs.tau] * vs.v_source
        vs.K_diss_bot = update_add(vs.K_diss_bot, at[...], numerics.calc_diss_v(state, diss))

    return KernelOutput(du_mix=vs.du_mix, dv_mix=vs.dv_mix, K_diss_bot=vs.K_diss_bot)


@veros_routine
def friction(state):
    vs = state.variables
    settings = state.settings

    """
    vertical friction
    """
    vs.K_diss_v = update(vs.K_diss_v, at[...], 0.0)

    if settings.enable_implicit_vert_friction:
        vs.update(implicit_vert_friction(state))

    if settings.enable_explicit_vert_friction:
        vs.update(explicit_vert_friction(state))

    """
    TEM formalism for eddy-driven velocity
    """
    if settings.enable_TEM_friction:
        vs.update(isoneutral.isoneutral_friction(state))

    """
    horizontal friction
    """
    if settings.enable_hor_friction:
        vs.update(harmonic_friction(state))

    if settings.enable_biharmonic_friction:
        vs.update(biharmonic_friction(state))

    """
    Rayleigh and bottom friction
    """
    vs.K_diss_bot = update(vs.K_diss_bot, at[...], 0.0)

    if settings.enable_ray_friction:
        vs.update(rayleigh_friction(state))

    if settings.enable_bottom_friction:
        vs.update(linear_bottom_friction(state))

    if settings.enable_quadratic_bottom_friction:
        vs.update(quadratic_bottom_friction(state))

    """
    add user defined forcing
    """
    if settings.enable_momentum_sources:
        vs.update(momentum_sources(state))
