from collections import OrderedDict

from . import veros_method


class Variable:
    def __init__(self, name, dims, units, long_description, dtype=None,
                 output=False, time_dependent=True, scale=1.,
                 write_to_restart=False, extra_attributes=None):
        self.name = name
        self.dims = dims
        self.units = units
        self.long_description = long_description
        self.dtype = dtype
        self.output = output
        self.time_dependent = time_dependent
        self.scale = scale
        self.write_to_restart = write_to_restart
        # : Additional attributes to be written in netCDF output
        self.extra_attributes = extra_attributes or {}


# fill value for netCDF output (invalid data is replaced by this value)
FILL_VALUE = -1e18

#
XT = ("xt",)
XU = ("xu",)
YT = ("yt",)
YU = ("yu",)
ZT = ("zt",)
ZW = ("zw",)
T_HOR = ("xt", "yt")
U_HOR = ("xu", "yt")
V_HOR = ("xt", "yu")
ZETA_HOR = ("xu", "yu")
T_GRID = ("xt", "yt", "zt")
U_GRID = ("xu", "yt", "zt")
V_GRID = ("xt", "yu", "zt")
W_GRID = ("xt", "yt", "zw")
ZETA_GRID = ("xu", "yu", "zt")
TIMESTEPS = ("timesteps",)
ISLE = ("isle",)
TENSOR_COMP = ("tensor1", "tensor2")

# those are written to netCDF output by default
BASE_DIMENSIONS = XT + XU + YT + YU + ZT + ZW + ISLE
GHOST_DIMENSIONS = ("xt", "yt", "xu", "yu")


def get_dimensions(vs, grid, include_ghosts=True):
    dimensions = {
        "xt": vs.nx,
        "xu": vs.nx,
        "yt": vs.ny,
        "yu": vs.ny,
        "zt": vs.nz,
        "zw": vs.nz,
        "timesteps": 3,
        "tensor1": 2,
        "tensor2": 2,
        "isle": vs.nisle,
    }
    if include_ghosts:
        for d in GHOST_DIMENSIONS:
            dimensions[d] += 4
    return tuple(dimensions[grid_dim] for grid_dim in grid)


def remove_ghosts(array, dims):
    ghost_mask = tuple(slice(2, -2) if dim in GHOST_DIMENSIONS else slice(None) for dim in dims)
    return array[ghost_mask]


@veros_method
def add_ghosts(vs, array, dims):
    full_shape = tuple([i + 4 if dim in GHOST_DIMENSIONS else i for i,
                        dim in zip(array.shape, dims)])
    newarr = np.zeros(full_shape, dtype=array.dtype)
    ghost_mask = tuple(slice(2, -2) if dim in GHOST_DIMENSIONS else slice(None) for dim in dims)
    newarr[ghost_mask] = array
    return newarr


def get_grid_mask(vs, grid):
    masks = {
        T_HOR: vs.maskT[:, :, -1],
        U_HOR: vs.maskU[:, :, -1],
        V_HOR: vs.maskV[:, :, -1],
        ZETA_HOR: vs.maskZ[:, :, -1],
        T_GRID: vs.maskT,
        U_GRID: vs.maskU,
        V_GRID: vs.maskV,
        W_GRID: vs.maskW,
        ZETA_GRID: vs.maskZ
    }
    if len(grid) > 2:
        if grid[:3] in masks.keys():
            return masks[grid[:3]]
    if len(grid) > 1:
        if grid[:2] in masks.keys():
            return masks[grid[:2]]
    return None


MAIN_VARIABLES = OrderedDict([
    ("dxt", Variable(
        "Zonal T-grid spacing", XT, "m",
        "Zonal (x) spacing of T-grid point",
        output=True, time_dependent=False
    )),
    ("dxu", Variable(
        "Zonal U-grid spacing", XU, "m",
        "Zonal (x) spacing of U-grid point",
        output=True, time_dependent=False
    )),
    ("dyt", Variable(
        "Meridional T-grid spacing", YT, "m",
        "Meridional (y) spacing of T-grid point",
        output=True, time_dependent=False
    )),
    ("dyu", Variable(
        "Meridional U-grid spacing", YU, "m",
        "Meridional (y) spacing of U-grid point",
        output=True, time_dependent=False
    )),
    ("zt", Variable(
        "Vertical coordinate (T)", ZT, "m", "Vertical coordinate",
        output=True, time_dependent=False, extra_attributes={"positive": "up"}
    )),
    ("zw", Variable(
        "Vertical coordinate (W)", ZW, "m", "Vertical coordinate", output=True,
        time_dependent=False, extra_attributes={"positive": "up"}
    )),
    ("dzt", Variable(
        "Vertical spacing (T)", ZT, "m", "Vertical spacing", output=True, time_dependent=False
    )),
    ("dzw", Variable(
        "Vertical spacing (W)", ZW, "m", "Vertical spacing", output=True, time_dependent=False
    )),
    ("cost", Variable(
        "Metric factor (T)", YT, "1", "Metric factor for spherical coordinates",
        time_dependent=False
    )),
    ("cosu", Variable(
        "Metric factor (U)", YU, "1", "Metric factor for spherical coordinates",
        time_dependent=False
    )),
    ("tantr", Variable(
        "Metric factor", YT, "1", "Metric factor for spherical coordinates",
        time_dependent=False
    )),
    ("coriolis_t", Variable(
        "Coriolis frequency", T_HOR, "1/s",
        "Coriolis frequency at T grid point", time_dependent=False
    )),
    ("coriolis_h", Variable(
        "Horizontal Coriolis frequency", T_HOR, "1/s",
        "Horizontal Coriolis frequency at T grid point", time_dependent=False
    )),

    ("kbot", Variable(
        "Index of deepest cell", T_HOR, "",
        "Index of the deepest grid cell (counting from 1, 0 means all land)",
        dtype="int", time_dependent=False
    )),
    ("ht", Variable(
        "Total depth (T)", T_HOR, "m", "Total depth of the water column", output=True,
        time_dependent=False
    )),
    ("hu", Variable(
        "Total depth (U)", U_HOR, "m", "Total depth of the water column", output=True,
        time_dependent=False
    )),
    ("hv", Variable(
        "Total depth (V)", V_HOR, "m", "Total depth of the water column", output=True,
        time_dependent=False
    )),
    ("hur", Variable(
        "Total depth (U), masked", U_HOR, "m",
        "Total depth of the water column (masked)", time_dependent=False
    )),
    ("hvr", Variable(
        "Total depth (V), masked", V_HOR, "m",
        "Total depth of the water column (masked)", time_dependent=False
    )),
    ("beta", Variable(
        "Change of Coriolis freq.", T_HOR, "1/(ms)",
        "Change of Coriolis frequency with latitude", output=True, time_dependent=False
    )),
    ("area_t", Variable(
        "Area of T-box", T_HOR, "m^2", "Area of T-box", output=True, time_dependent=False
    )),
    ("area_u", Variable(
        "Area of U-box", U_HOR, "m^2", "Area of U-box", output=True, time_dependent=False
    )),
    ("area_v", Variable(
        "Area of V-box", V_HOR, "m^2", "Area of V-box", output=True, time_dependent=False
    )),

    ("maskT", Variable(
        "Mask for tracer points", T_GRID, "",
        "Mask in physical space for tracer points", dtype="int", time_dependent=False
    )),
    ("maskU", Variable(
        "Mask for U points", U_GRID, "",
        "Mask in physical space for U points", dtype="int", time_dependent=False
    )),
    ("maskV", Variable(
        "Mask for V points", V_GRID, "",
        "Mask in physical space for V points", dtype="int", time_dependent=False
    )),
    ("maskW", Variable(
        "Mask for W points", W_GRID, "",
        "Mask in physical space for W points", dtype="int", time_dependent=False
    )),
    ("maskZ", Variable(
        "Mask for Zeta points", ZETA_GRID, "",
        "Mask in physical space for Zeta points", dtype="int", time_dependent=False
    )),

    ("rho", Variable(
        "Density", T_GRID + TIMESTEPS, "kg/m^3", "Potential density", output=True,
        write_to_restart=True
    )),
    ("int_drhodT", Variable(
        "Der. of dyn. enthalpy by temperature", T_GRID + TIMESTEPS, "?",
        "Partial derivative of dynamic enthalpy by temperature", output=True,
        write_to_restart=True
    )),
    ("int_drhodS", Variable(
        "Der. of dyn. enthalpy by salinity", T_GRID + TIMESTEPS, "?",
        "Partial derivative of dynamic enthalpy by salinity", output=True,
        write_to_restart=True
    )),
    ("Nsqr", Variable(
        "Square of stability frequency", W_GRID + TIMESTEPS, "1/s^2",
        "Square of stability frequency", output=True, write_to_restart=True
    )),
    ("Hd", Variable(
        "Dynamic enthalpy", T_GRID + TIMESTEPS, "m^2/s^2", "Dynamic enthalpy",
        output=True, write_to_restart=True
    )),
    ("dHd", Variable(
        "Change of dyn. enth. by adv.", T_GRID + TIMESTEPS, "m^2/s^3",
        "Change of dynamic enthalpy due to advection", write_to_restart=True
    )),

    ("temp", Variable(
        "Temperature", T_GRID + TIMESTEPS, "deg C",
        "Conservative temperature", output=True, write_to_restart=True
    )),
    ("dtemp", Variable(
        "Temperature tendency", T_GRID + TIMESTEPS, "deg C/s",
        "Conservative temperature tendency", write_to_restart=True
    )),
    ("salt", Variable(
        "Salinity", T_GRID + TIMESTEPS, "g/kg", "Salinity", output=True,
        write_to_restart=True
    )),
    ("dsalt", Variable(
        "Salinity tendency", T_GRID + TIMESTEPS, "g/(kg s)",
        "Salinity tendency", write_to_restart=True
    )),
    ("dtemp_vmix", Variable(
        "Change of temp. by vertical mixing", T_GRID, "deg C/s",
        "Change of temperature due to vertical mixing",
    )),
    ("dtemp_hmix", Variable(
        "Change of temp. by horizontal mixing", T_GRID, "deg C/s",
        "Change of temperature due to horizontal mixing",
    )),
    ("dsalt_vmix", Variable(
        "Change of sal. by vertical mixing", T_GRID, "deg C/s",
        "Change of salinity due to vertical mixing",
    )),
    ("dsalt_hmix", Variable(
        "Change of sal. by horizontal mixing", T_GRID, "deg C/s",
        "Change of salinity due to horizontal mixing",
    )),
    ("dtemp_iso", Variable(
        "Change of temp. by isop. mixing", T_GRID, "deg C/s",
        "Change of temperature due to isopycnal mixing plus skew mixing",
    )),
    ("dsalt_iso", Variable(
        "Change of sal. by isop. mixing", T_GRID, "deg C/s",
        "Change of salinity due to isopycnal mixing plus skew mixing",

    )),
    ("forc_temp_surface", Variable(
        "Surface temperature flux", T_HOR, "m K/s", "Surface temperature flux",
        output=True
    )),
    ("forc_salt_surface", Variable(
        "Surface salinity flux", T_HOR, "m g/s kg", "Surface salinity flux",
        output=True
    )),

    ("flux_east", Variable(
        "Multi-purpose flux", U_GRID, "?", "Multi-purpose flux"
    )),
    ("flux_north", Variable(
        "Multi-purpose flux", V_GRID, "?", "Multi-purpose flux"
    )),
    ("flux_top", Variable(
        "Multi-purpose flux", W_GRID, "?", "Multi-purpose flux"
    )),

    ("u", Variable(
        "Zonal velocity", U_GRID + TIMESTEPS, "m/s", "Zonal velocity",
        output=True, write_to_restart=True
    )),
    ("v", Variable(
        "Meridional velocity", V_GRID + TIMESTEPS, "m/s", "Meridional velocity",
        output=True, write_to_restart=True
    )),
    ("w", Variable(
        "Vertical velocity", W_GRID + TIMESTEPS, "m/s", "Vertical velocity",
        output=True, write_to_restart=True
    )),
    ("du", Variable(
        "Zonal velocity tendency", U_GRID + TIMESTEPS, "m/s",
        "Zonal velocity tendency", write_to_restart=True
    )),
    ("dv", Variable(
        "Meridional velocity tendency", V_GRID + TIMESTEPS, "m/s",
        "Meridional velocity tendency", write_to_restart=True
    )),
    ("du_cor", Variable(
        "Change of u by Coriolis force", U_GRID, "m/s^2",
        "Change of u due to Coriolis force"
    )),
    ("dv_cor", Variable(
        "Change of v by Coriolis force", V_GRID, "m/s^2",
        "Change of v due to Coriolis force"
    )),
    ("du_mix", Variable(
        "Change of u by vertical mixing", U_GRID, "m/s^2",
        "Change of u due to implicit vertical mixing"
    )),
    ("dv_mix", Variable(
        "Change of v by vertical mixing", V_GRID, "m/s^2",
        "Change of v due to implicit vertical mixing"
    )),
    ("du_adv", Variable(
        "Change of u by advection", U_GRID, "m/s^2",
        "Change of u due to advection"
    )),
    ("dv_adv", Variable(
        "Change of v by advection", V_GRID, "m/s^2",
        "Change of v due to advection"
    )),
    ("p_hydro", Variable(
        "Hydrostatic pressure", T_GRID, "m^2/s^2", "Hydrostatic pressure", output=True
    )),
    ("kappaM", Variable(
        "Vertical viscosity", T_GRID, "m^2/s", "Vertical viscosity", output=True
    )),
    ("kappaH", Variable(
        "Vertical diffusivity", W_GRID, "m^2/s", "Vertical diffusivity", output=True
    )),
    ("surface_taux", Variable(
        "Surface wind stress", U_HOR, "m^2/s^2", "Zonal surface wind stress", output=True
    )),
    ("surface_tauy", Variable(
        "Surface wind stress", V_HOR, "m^2/s^2", "Meridional surface wind stress", output=True
    )),
    ("forc_rho_surface", Variable(
        "Surface density flux", T_HOR, "?", "Surface potential density flux", output=True
    )),

    ("psi", Variable(
        "Streamfunction", ZETA_HOR + TIMESTEPS, "m^3/s", "Streamfunction",
        output=True, write_to_restart=True
    )),
    ("dpsi", Variable(
        "Streamfunction tendency", ZETA_HOR + TIMESTEPS, "m^3/s^2",
        "Streamfunction tendency", write_to_restart=True
    )),
    ("isle", Variable(
        "Island number", ISLE, "", "Island number", output=True
    )),
    ("psin", Variable(
        "Boundary streamfunction", ZETA_HOR + ISLE, "m^3/s",
        "Boundary streamfunction", output=True, time_dependent=False
    )),
    ("dpsin", Variable(
        "Boundary streamfunction factor", ISLE + TIMESTEPS, "?",
        "Boundary streamfunction factor", write_to_restart=True
    )),
    ("line_psin", Variable(
        "Boundary line integrals", ISLE + ISLE, "?",
        "Boundary line integrals", time_dependent=False
    )),
    ("boundary_mask", Variable(
        "Boundary mask", T_HOR + ISLE, "",
        "Boundary mask", time_dependent=False
    )),
    ("line_dir_south_mask", Variable(
        "Line integral mask", T_HOR + ISLE, "",
        "Line integral mask", time_dependent=False
    )),
    ("line_dir_north_mask", Variable(
        "Line integral mask", T_HOR + ISLE, "",
        "Line integral mask", time_dependent=False
    )),
    ("line_dir_east_mask", Variable(
        "Line integral mask", T_HOR + ISLE, "",
        "Line integral mask", time_dependent=False
    )),
    ("line_dir_west_mask", Variable(
        "Line integral mask", T_HOR + ISLE, "",
        "Line integral mask", time_dependent=False
    )),

    ("K_gm", Variable(
        "Skewness diffusivity", W_GRID, "m^2/s",
        "GM diffusivity, either constant or from EKE model"
    )),
    ("K_iso", Variable(
        "Isopycnal diffusivity", W_GRID, "m^2/s", "Along-isopycnal diffusivity"
    )),

    ("K_diss_v", Variable(
        "Dissipation of kinetic Energy", W_GRID, "m^2/s^3",
        "Kinetic energy dissipation by vertical, rayleigh and bottom friction",
        write_to_restart=True
    )),
    ("K_diss_bot", Variable(
        "Dissipation of kinetic Energy", W_GRID, "m^2/s^3",
        "Mean energy dissipation by bottom and rayleigh friction"
    )),
    ("K_diss_h", Variable(
        "Dissipation of kinetic Energy", W_GRID, "m^2/s^3",
        "Kinetic energy dissipation by horizontal friction"
    )),
    ("K_diss_gm", Variable(
        "Dissipation of mean energy", W_GRID, "m^2/s^3",
        "Mean energy dissipation by GM (TRM formalism only)"
    )),
    ("P_diss_v", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by vertical diffusion"
    )),
    ("P_diss_nonlin", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by nonlinear equation of state"
    )),
    ("P_diss_iso", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by isopycnal mixing"
    )),
    ("P_diss_skew", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by GM (w/o TRM)"
    )),
    ("P_diss_hmix", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by horizontal mixing"
    )),
    ("P_diss_adv", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by advection"
    )),
    ("P_diss_comp", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by compression"
    )),
    ("P_diss_sources", Variable(
        "Dissipation of potential Energy", W_GRID, "m^2/s^3",
        "Potential energy dissipation by external sources (e.g. restoring zones)"
    )),

    ("u_wgrid", Variable(
        "U on W grid", W_GRID, "m/s", "Zonal velocity interpolated to W grid points"
    )),
    ("v_wgrid", Variable(
        "V on W grid", W_GRID, "m/s", "Meridional velocity interpolated to W grid points"
    )),
    ("w_wgrid", Variable(
        "W on W grid", W_GRID, "m/s", "Vertical velocity interpolated to W grid points"
    ))
])

CONDITIONAL_VARIABLES = OrderedDict([
    ("coord_degree", OrderedDict([
        ("xt", Variable(
            "Zonal coordinate (T)", XT, "degrees_east",
            "Zonal (x) coordinate of T-grid point",
            output=True, time_dependent=False
        )),
        ("xu", Variable(
            "Zonal coordinate (U)", XU, "degrees_east",
            "Zonal (x) coordinate of U-grid point",
            output=True, time_dependent=False
        )),
        ("yt", Variable(
            "Meridional coordinate (T)", YT, "degrees_north",
            "Meridional (y) coordinate of T-grid point",
            output=True, time_dependent=False
        )),
        ("yu", Variable(
            "Meridional coordinate (U)", YU, "degrees_north",
            "Meridional (y) coordinate of U-grid point",
            output=True, time_dependent=False
        )),
    ])),

    ("not coord_degree", OrderedDict([
        ("xt", Variable(
            "Zonal coordinate (T)", XT, "km",
            "Zonal (x) coordinate of T-grid point",
            output=True, scale=1e-3, time_dependent=False
        )),
        ("xu", Variable(
            "Zonal coordinate (U)", XU, "km",
            "Zonal (x) coordinate of U-grid point",
            output=True, scale=1e-3, time_dependent=False
        )),
        ("yt", Variable(
            "Meridional coordinate (T)", YT, "km",
            "Meridional (y) coordinate of T-grid point",
            output=True, scale=1e-3, time_dependent=False
        )),
        ("yu", Variable(
            "Meridional coordinate (U)", YU, "km",
            "Meridional (y) coordinate of U-grid point",
            output=True, scale=1e-3, time_dependent=False
        )),
    ])),

    ("enable_tempsalt_sources", OrderedDict([
        ("temp_source", Variable(
            "Source of temperature", T_GRID, "K/s",
            "Non-conservative source of temperature", output=True
        )),
        ("salt_source", Variable(
            "Source of salt", T_GRID, "g/(kg s)",
            "Non-conservative source of salt", output=True
        )),
    ])),

    ("enable_momentum_sources", OrderedDict([
        ("u_source", Variable(
            "Source of zonal velocity", U_GRID, "m/s^2 (?)",
            "Non-conservative source of zonal velocity", output=True
        )),
        ("v_source", Variable(
            "Source of meridional velocity", V_GRID, "m/s^2 (?)",
            "Non-conservative source of meridional velocity", output=True
        )),
    ])),

    ("enable_neutral_diffusion", OrderedDict([
        ("K_11", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_13", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_22", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_23", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_31", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_32", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("K_33", Variable("Isopycnal mixing coefficient", T_GRID, "?", "Isopycnal mixing tensor component")),
        ("Ai_ez", Variable("?", T_GRID + TENSOR_COMP, "?", "?")),
        ("Ai_nz", Variable("?", T_GRID + TENSOR_COMP, "?", "?")),
        ("Ai_bx", Variable("?", T_GRID + TENSOR_COMP, "?", "?")),
        ("Ai_by", Variable("?", T_GRID + TENSOR_COMP, "?", "?")),
    ])),
    ("enable_skew_diffusion", OrderedDict([
        ("B1_gm", Variable(
            "Zonal component of GM streamfunction", V_GRID, "m^2/s",
            "Zonal component of GM streamfunction"
        )),
        ("B2_gm", Variable(
            "Meridional component of GM streamfunction", U_GRID, "m^2/s",
            "Meridional component of GM streamfunction"
        ))
    ])),
    ("enable_bottom_friction_var", OrderedDict([
        ("r_bot_var_u", Variable(
            "Bottom friction coeff.", U_HOR, "?", "Zonal bottom friction coefficient"
        )),
        ("r_bot_var_v", Variable(
            "Bottom friction coeff.", V_HOR, "?", "Meridional bottom friction coefficient"
        )),
    ])),
    ("enable_TEM_friction", OrderedDict([
        ("kappa_gm", Variable("Vertical diffusivity", W_GRID, "m^2/s", "Vertical diffusivity")),
    ])),
    ("enable_tke", OrderedDict([
        ("tke", Variable(
            "Turbulent kinetic energy", W_GRID + TIMESTEPS, "m^2/s^2",
            "Turbulent kinetic energy", output=True, write_to_restart=True
        )),
        ("sqrttke", Variable(
            "Square-root of TKE", W_GRID, "m/s", "Square-root of TKE"
        )),
        ("dtke", Variable(
            "Turbulent kinetic energy tendency", W_GRID + TIMESTEPS, "m^2/s^3",
            "Turbulent kinetic energy tendency", write_to_restart=True
        )),
        ("Prandtlnumber", Variable("Prandtl number", W_GRID, "", "Prandtl number")),
        ("mxl", Variable("Mixing length", W_GRID, "m", "Mixing length")),
        ("forc_tke_surface", Variable(
            "TKE surface flux", T_HOR, "m^3/s^3", "TKE surface flux", output=True
        )),
        ("tke_diss", Variable(
            "TKE dissipation", W_GRID, "m^2/s^3", "TKE dissipation"
        )),
        ("tke_surf_corr", Variable(
            "Correction of TKE surface flux", T_HOR, "m^3/s^3",
            "Correction of TKE surface flux"
        )),
    ])),
    ("enable_eke", OrderedDict([
        ("eke", Variable(
            "meso-scale energy", W_GRID + TIMESTEPS, "m^2/s^2",
            "meso-scale energy", output=True, write_to_restart=True
        )),
        ("deke", Variable(
            "meso-scale energy tendency", W_GRID + TIMESTEPS, "m^2/s^3",
            "meso-scale energy tendency", write_to_restart=True
        )),
        ("sqrteke", Variable(
            "square-root of eke", W_GRID, "m/s", "square-root of eke"
        )),
        ("L_rossby", Variable("Rossby radius", T_HOR, "m", "Rossby radius")),
        ("L_rhines", Variable("Rhines scale", W_GRID, "m", "Rhines scale")),
        ("eke_len", Variable("Eddy length scale", T_GRID, "m", "Eddy length scale")),
        ("eke_diss_iw", Variable(
            "Dissipation of EKE to IW", W_GRID, "m^2/s^3",
            "Dissipation of EKE to internal waves"
        )),
        ("eke_diss_tke", Variable(
            "Dissipation of EKE to TKE", W_GRID, "m^2/s^3",
            "Dissipation of EKE to TKE"
        )),
        ("eke_bot_flux", Variable(
            "Flux by bottom friction", T_HOR, "m^3/s^3", "Flux by bottom friction"
        )),
    ])),
    ("enable_eke_leewave_dissipation", OrderedDict([
        ("eke_topo_hrms", Variable(
            "?", T_HOR, "?", "?"
        )),
        ("eke_topo_lam", Variable(
            "?", T_HOR, "?", "?"
        )),
        ("hrms_k0", Variable(
            "?", T_HOR, "?", "?"
        )),
        ("c_lee", Variable(
            "Lee wave dissipation coefficient", T_HOR, "1/s",
            "Lee wave dissipation coefficient"
        )),
        ("eke_lee_flux", Variable(
            "Lee wave flux", T_HOR, "m^3/s^3", "Lee wave flux",
        )),
        ("c_Ri_diss", Variable(
            "Interior dissipation coefficient", W_GRID, "1/s",
            "Interior dissipation coefficient"
        )),
    ])),
    ("enable_idemix", OrderedDict([
        ("E_iw", Variable(
            "Internal wave energy", W_GRID + TIMESTEPS, "m^2/s^2",
            "Internal wave energy", output=True, write_to_restart=True
        )),
        ("dE_iw", Variable(
            "Internal wave energy tendency", W_GRID + TIMESTEPS, "m^2/s^2",
            "Internal wave energy tendency", write_to_restart=True
        )),
        ("c0", Variable(
            "Vertical IW group velocity", W_GRID, "m/s",
            "Vertical internal wave group velocity"
        )),
        ("v0", Variable(
            "Horizontal IW group velocity", W_GRID, "m/s",
            "Horizontal internal wave group velocity"
        )),
        ("alpha_c", Variable("?", W_GRID, "?", "?")),
        ("iw_diss", Variable(
            "IW dissipation", W_GRID, "m^2/s^3", "Internal wave dissipation"
        )),
        ("forc_iw_surface", Variable(
            "IW surface forcing", T_HOR, "m^3/s^3",
            "Internal wave surface forcing", time_dependent=False, output=True
        )),
        ("forc_iw_bottom", Variable(
            "IW bottom forcing", T_HOR, "m^3/s^3",
            "Internal wave bottom forcing", time_dependent=False, output=True
        )),
    ])),
])


@veros_method
def allocate_variables(vs):
    variables = {}

    def init_var(var_name, var):
        shape = get_dimensions(vs, var.dims)

        kwargs = {}
        kwargs["dtype"] = var.dtype or vs.default_float_type
        if vs.backend_name == "dask":
            kwargs["chunks"] = shape

        setattr(vs, var_name, np.zeros(shape, **kwargs))
        variables[var_name] = var

    for var_name, var in MAIN_VARIABLES.items():
        init_var(var_name, var)

    for condition, var_dict in CONDITIONAL_VARIABLES.items():
        if condition.startswith("not "):
            eval_condition = not bool(getattr(vs, condition[4:]))
        else:
            eval_condition = bool(getattr(vs, condition))
        if eval_condition:
            for var_name, var in var_dict.items():
                init_var(var_name, var)

    return variables
