"""
Common constants giving information about available IO and their meta-data.

Constants
=========

.. autosummary::

    OUTPUT_META
    OUTPUT_HORIZONS_META
    OUTPUT_CALC_META
    OUTPUT_CALC_HORIZONS_META
    MRM_OUTPUT_META
    INPUT_UNITS

----

.. autodata:: OUTPUT_META
    :no-value:
.. pprint:: OUTPUT_META

.. autodata:: OUTPUT_HORIZONS_META
    :no-value:
.. pprint:: OUTPUT_HORIZONS_META

.. autodata:: OUTPUT_CALC_META
    :no-value:
.. pprint:: OUTPUT_CALC_META

.. autodata:: OUTPUT_CALC_HORIZONS_META
    :no-value:
.. pprint:: OUTPUT_CALC_HORIZONS_META

.. autodata:: MRM_OUTPUT_META
    :no-value:
.. pprint:: MRM_OUTPUT_META

.. autodata:: INPUT_UNITS
    :no-value:
.. pprint:: INPUT_UNITS

"""

# pylint: disable=R1735
import mhm
import numpy as np

OUTPUT_META = {
    "L0_GRIDDED_LAI": dict(units="1", long_name="leaf area index"),
    "L1_FSEALED": dict(units="1", long_name="Fraction of sealed area"),
    "L1_FNOTSEALED": dict(units="1", long_name="Fraction of unsealed area"),
    "L1_INTER": dict(units="mm / h", long_name="Interception"),  # interception (1)
    "L1_SNOWPACK": dict(units="mm / h", long_name="Snowpack"),  # snowpack (2)
    "L1_SEALSTW": dict(
        units="mm / h", long_name="Retention storage of impervious areas"
    ),  # sealedSTW (6)
    "L1_UNSATSTW": dict(units="mm / h", long_name="upper soil storage"),  # unsatSTW (7)
    "L1_SATSTW": dict(units="mm / h", long_name="groundwater storage"),  # satSTW (8)
    "L1_NEUTRONS": dict(
        units="mm / h", long_name="Ground Albedo Neutrons"
    ),  # neutrons (18)
    "L1_PET_CALC": dict(
        units="mm / h", long_name="potential evapotranspiration"
    ),  # PET (9)
    "L1_AETCANOPY": dict(
        units="mm / h", long_name="Real evaporation intensity from canopy"
    ),
    "L1_AETSEALED": dict(
        units="mm / h", long_name="Real evap. from free water surfaces"
    ),
    "L1_TOTAL_RUNOFF": dict(units="mm / h", long_name="Generated runoff"),  # Q (11)
    "L1_RUNOFFSEAL": dict(
        units="mm / h", long_name="Direct runoff from impervious areas"
    ),
    "L1_FASTRUNOFF": dict(units="mm / h", long_name="Fast runoff component"),
    "L1_SLOWRUNOFF": dict(units="mm / h", long_name="Slow runoff component"),
    "L1_BASEFLOW": dict(units="mm / h", long_name="Baseflow"),
    "L1_PERCOL": dict(units="mm / h", long_name="Percolation"),
    "L1_PREEFFECT": dict(
        units="mm / h", long_name="Effective precip. depth"
    ),  # preEffect (20)
    "L1_SOILMOIST_VOL_ALL": dict(
        units="1", long_name="average soil moisture over all layers"
    ),  # SM_Lall (5)
}
"""meta information about available outputs in mHM."""

MRM_OUTPUT_META = {
    "L11_QMOD": dict(units="m^3 / s", long_name="Simulated discharge"),
    "L11_QOUT": dict(units="m^3 / s", long_name="Total outflow from cells"),
}
"""meta information about available outputs in mRM."""

OUTPUT_HORIZONS_META = {
    "L1_SOILMOIST": dict(
        units="mm", long_name="soil water content of soil layer {n}"
    ),  # SWC_Lxx (3)
    "L1_SOILMOIST_VOL": dict(
        units="mm / mm", long_name="volumetric soil moisture of soil layer {n}"
    ),  # SM_Lxx (4)
    "L1_SOILMOISTSAT": dict(
        units="mm", long_name="saturation soil moisture of soil layer {n}"
    ),
    "L1_AETSOIL": dict(units="mm / h", long_name="actual ET  of soil layer {n}"),
    "L1_INFILSOIL": dict(
        units="mm / h", long_name="infiltration intensity of soil layer {n}"
    ),
    # missing
    # L1_MELT
    # L1_RAIN
    # L1_SNOW
    # L1_THROUGHFALL
}
"""meta information about available outputs per horizon in mHM."""

OUTPUT_CALC_META = {
    # sum(aETSoil(horizons)) * fNotSealed + aETCanopy + aETSealed * fSealed
    "L1_AET": dict(units="mm / h", long_name="actual Evapotranspiration"),  # aET (10)
    # runoffSeal * fSealed
    "L1_QD": dict(
        units="mm / h", long_name="direct runoff generated by every cell (runoffSeal)"
    ),  # QD (12)
    # fastRunoff * fNotSealed
    "L1_QIF": dict(
        units="mm / h", long_name="fast interflow generated by every cell (fastRunoff)"
    ),  # QIf (13)
    # slowRunoff * fNotSealed
    "L1_QIS": dict(
        units="mm / h", long_name="slow interflow generated by every cell (slowRunoff)"
    ),  # QIs (14)
    # baseflow * fNotSealed
    "L1_QB": dict(
        units="mm / h", long_name="baseflow generated by every cell"
    ),  # QB (15)
    # percol * fNotSealed
    "L1_RECHARGE": dict(
        units="mm / h", long_name="groundwater recharge"
    ),  # recharge (16)
}
"""meta information about calculated outputs in mHM."""

OUTPUT_CALC_HORIZONS_META = {
    # infilSoil(horizon) * fNotSealed
    "L1_SOIL_INFIL": dict(
        units="mm / h", long_name="infiltration flux from soil layer {n}"
    ),  # soil_infil_Lxx (17)
    # aETSoil(horizon) * fNotSealed
    "L1_AET": dict(
        units="mm / h", long_name="actual Evapotranspiration from soil layer {n}"
    ),  # aET_Lxx (19)
}
"""meta information about calculated outputs per horizon in mHM."""

INPUT_UNITS = {
    # "L0_GRIDDED_LAI": "1",
    "METEO_PRE": "mm / {ts}",
    "METEO_TEMP": "degC",
    "METEO_PET": "mm / {ts}",
    "METEO_TMIN": "degC",
    "METEO_TMAX": "degC",
    "METEO_NETRAD": "W m-2",
    "METEO_ABSVAPPRESS": "Pa",
    "METEO_WINDSPEED": "m s-1",
    "METEO_SSRD": "W m-2",
    "METEO_STRD": "W m-2",
    "METEO_TANN": "degC",
}
"""units of the available inputs in mHM."""

HOURS_TO_TIMESTEP = {1: "h", 24: "d"}
"""timestep string from hours."""


def _fill_var(var, grid="l1"):
    grid_info = getattr(mhm.get, grid + "_domain_info")()
    # mask in mHM is the opposite in numpy (selection)
    sel = mhm.get_mask(grid, indexing="xy", selection=True)
    sel = sel.ravel(order="F")
    output = np.ma.empty_like(sel, dtype=float)
    output.fill_value = grid_info[-1]
    output.mask = ~sel
    output[sel] = var
    return output.reshape((grid_info[0], grid_info[1]), order="C")


def _L1_AET():
    # sum(aETSoil(horizons)) * fNotSealed + aETCanopy + aETSealed * fSealed
    fsealed = mhm.get_variable("L1_FSEALED", compressed=True)
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    aetcanopy = mhm.get_variable("L1_AETCANOPY", compressed=True)
    aetsealed = mhm.get_variable("L1_AETSEALED", compressed=True)
    horizons = mhm.get.number_of_horizons()

    sum_aetsoil = np.zeros_like(fsealed, dtype=float)
    for n in range(1, horizons + 1):
        sum_aetsoil += mhm.get_variable("L1_AETSOIL", index=n, compressed=True)

    return _fill_var(sum_aetsoil * fnotsealed + aetcanopy + aetsealed * fsealed)


def _L1_QD():
    # runoffSeal * fSealed
    fsealed = mhm.get_variable("L1_FSEALED", compressed=True)
    runoffseal = mhm.get_variable("L1_RUNOFFSEAL", compressed=True)
    return _fill_var(runoffseal * fsealed)


def _L1_QIF():
    # fastRunoff * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    fastrunoff = mhm.get_variable("L1_FASTRUNOFF", compressed=True)
    return _fill_var(fastrunoff * fnotsealed)


def _L1_QIS():
    # slowRunoff * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    slowrunoff = mhm.get_variable("L1_SLOWRUNOFF", compressed=True)
    return _fill_var(slowrunoff * fnotsealed)


def _L1_QB():
    # baseflow * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    baseflow = mhm.get_variable("L1_BASEFLOW", compressed=True)
    return _fill_var(baseflow * fnotsealed)


def _L1_RECHARGE():
    # percol * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    percol = mhm.get_variable("L1_PERCOL", compressed=True)
    return _fill_var(percol * fnotsealed)


def _L1_SOIL_INFIL_N(n):
    # infilSoil(horizon) * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    infilsoil = mhm.get_variable("L1_INFILSOIL", index=n, compressed=True)
    return _fill_var(infilsoil * fnotsealed)


def _L1_AET_N(n):
    # aETSoil(horizon) * fNotSealed
    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
    aetsoil = mhm.get_variable("L1_AETSOIL", index=n, compressed=True)
    return _fill_var(aetsoil * fnotsealed)


OUTPUT_CALC = {
    "L1_AET": _L1_AET,
    "L1_QD": _L1_QD,
    "L1_QIF": _L1_QIF,
    "L1_QIS": _L1_QIS,
    "L1_QB": _L1_QB,
    "L1_RECHARGE": _L1_RECHARGE,
}

OUTPUT_CALC_HORIZON = {
    "L1_SOIL_INFIL": _L1_SOIL_INFIL_N,
    "L1_AET": _L1_AET_N,
}
