__init__.py

"""
Read MODFLOW6 output

The dis, disv, disu modules implement the following functions:

```python
Darray = Union[xr.DataArray, xu.UgridDataArray]

def read_grb(f: BinaryIO, ntxt: int, lentxt: int) -> Dict[str, Any]:
    return

def read_times(*args) -> FloatArray:
    return

def read_hds_timestep(*args) -> FloatArray:
    return

def open_hds(path: FilePath, d: Dict[str, Any], dry_nan: bool) -> Darray:
    return

def open_imeth1_budgets(
    cbc_path: FilePath, grb_content: dict, header_list: List["Imeth1Header"]
) -> Darray:
    return

def open_imeth6_budgets(
    cbc_path: FilePath, grb_content: dict, header_list: List["Imeth6Header"]
) -> Darray:
    return

def open_cbc(
    cbc_path: FilePath, grb_content: Dict[str, Any]
) -> Dict[str, Darray]:
    return
```

(These could be implemented via Reader classes, but why bother with mutable
state or a class with exclusively staticmethods?)
"""

from typing import Any, Callable, Dict, Optional, Union

import numpy as np
import xarray as xr
import xugrid as xu

from imod.typing import GridDataArray, GridDataset
from imod.typing.grid import merge_with_dictionary

from . import dis, disu, disv
from .cbc import read_cbc_headers
from .common import FilePath, _grb_text

_READ_GRB = {
    "grid dis": dis.read_grb,
    "grid disv": disv.read_grb,
    "grid disu": disu.read_grb,
}

_OPEN_HDS = {
    "dis": dis.open_hds,
    "disv": disv.open_hds,
    "disu": disu.open_hds,
}

_OPEN_CBC: Dict[str, Callable[..., Any]] = {
    "dis": dis.open_cbc,
    "disv": disv.open_cbc,
    "disu": disu.open_cbc,
}

_OPEN_DVS = {
    "dis": dis.open_dvs,
    "disv": disv.open_dvs,
    "disu": disu.open_dvs,
}


def _get_function(d: Dict[str, Callable[..., Any]], key: str) -> Callable[..., Any]:
    try:
        func = d[key]
    except KeyError:
        valid_options = ", ".join(d.keys()).lower()
        raise ValueError(f"Expected one of {valid_options}, got: {key}")
    return func


def read_grb(path: FilePath) -> Dict[str, Any]:
    """
    Read the data in a MODFLOW6 binary grid (.grb) file.

    Parameters
    ----------
    path: Union[str, pathlib.Path]

    Returns
    -------
    grb_content: Dict[str, Any]

    Examples
    --------
    Read the grb file:

    >>> grb_content = imod.mf6.read_grb("my-model.grb")
    """
    with open(path, "rb") as f:
        h1 = _grb_text(f)
        _read = _get_function(_READ_GRB, h1)
        h2 = _grb_text(f)
        if h2 != "version 1":
            raise ValueError(f"Only version 1 supported, got {h2}")
        ntxt = int(_grb_text(f).split()[1])
        lentxt = int(_grb_text(f).split()[1])
        d = _read(f, ntxt, lentxt)
    return d


def open_hds(
    hds_path: FilePath,
    grb_path: FilePath,
    dry_nan: bool = False,
    simulation_start_time: Optional[np.datetime64] = None,
    time_unit: Optional[str] = "d",
) -> GridDataArray:
    """
    Open modflow6 heads (.hds) file.

    The data is lazily read per timestep and automatically converted into
    (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV respectively.
    The conversion is done via the information stored in the Binary Grid file
    (GRB).


    Parameters
    ----------
    hds_path: Union[str, pathlib.Path]
    grb_path: Union[str, pathlib.Path]
    dry_nan: bool, default value: False.
        Whether to convert dry values to NaN.
    simulation_start_time : Optional datetime
        The time and date correpsonding to the beginning of the simulation.
        Use this to convert the time coordinates of the output array to
        calendar time/dates. time_unit must also be present if this argument is present.
    time_unit: Optional str
        The time unit MF6 is working in, in string representation.
        Only used if simulation_start_time was provided.
        Admissible values are:
        ns -> nanosecond
        ms -> microsecond
        s -> second
        m -> minute
        h -> hour
        d -> day
        w -> week
        Units "month" or "year" are not supported, as they do not represent unambiguous timedelta values durations.

    Returns
    -------
    head: Union[xr.DataArray, xu.UgridDataArray]

    Examples
    --------
    When you completed MODFLOW6 simulation, you can open the head file:

    >>> head = imod.mf6.open_hds("my-model.hds", "my-model.grb")

    This will return a DataArray with the head values for each cell in the
    model, with dimensions ("time", "layer", "y", "x"). The time coordinate will
    only be the incremental time steps (1.0, 2.0, ...), unless you provide the
    ``simulation_start_time`` and ``time_unit`` arguments. This will convert the
    time coordinate to calendar time, starting from the provided simulation
    start time:

    >>> head = imod.mf6.open_hds(
    ...     "my-model.hds",
    ...     "my-model.grb",
    ...     simulation_start_time=np.datetime64("2023-01-01T00:00:00"),
    ...     time_unit="d",
    ... )

    """
    grb_content = read_grb(grb_path)
    grb_content["name"] = "head"
    distype = grb_content["distype"]
    _open = _get_function(_OPEN_HDS, distype)
    return _open(hds_path, grb_content, dry_nan, simulation_start_time, time_unit)


def open_dvs(
    dvs_path: FilePath,
    grb_path: FilePath,
    indices: np.ndarray,
    simulation_start_time: Optional[np.datetime64] = None,
    time_unit: Optional[str] = "d",
) -> GridDataArray:
    """
    Open modflow6 dependent variable output for complex packages (like the watercontent file for the UZF-package).

    The data is lazily read per timestep and automatically converted into
    (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV respectively.
    The conversion is done via the information stored in the Binary Grid file
    (GRB).


    Parameters
    ----------
    dvs_path: Union[str, pathlib.Path]
    grb_path: Union[str, pathlib.Path]
    indices: np.ndarray
        The indices to map dvs variable to model nodes
    simulation_start_time : Optional datetime
        The time and date correpsonding to the beginning of the simulation.
        Use this to convert the time coordinates of the output array to
        calendar time/dates. time_unit must also be present if this argument is present.
    time_unit: Optional str
        The time unit MF6 is working in, in string representation.
        Only used if simulation_start_time was provided.
        Admissible values are:
        ns -> nanosecond
        ms -> microsecond
        s -> second
        m -> minute
        h -> hour
        d -> day
        w -> week
        Units "month" or "year" are not supported, as they do not represent unambiguous timedelta values durations.

    Returns
    -------
    dvs : Union[xr.DataArray, xu.UgridDataArray]

    Examples
    --------
    When you completed MODFLOW6 simulation, you can open the dvs file:

    >>> dvs = imod.mf6.open_dvs("my-model.dvs", "my-model.grb")

    This will return a DataArray with the dvs values for each cell in the
    model, with dimensions ("time", "layer", "y", "x"). The time coordinate will
    only be the incremental time steps (1.0, 2.0, ...), unless you provide the
    ``simulation_start_time`` and ``time_unit`` arguments. This will convert the
    time coordinate to calendar time, starting from the provided simulation
    start time:

    >>> dvs = imod.mf6.open_dvs(
    ...     "my-model.dvs",
    ...     "my-model.grb",
    ...     simulation_start_time=np.datetime64("2023-01-01T00:00:00"),
    ...     time_unit="d",
    ... )

    """
    grb_content = read_grb(grb_path)
    distype = grb_content["distype"]
    _open = _get_function(_OPEN_DVS, distype)
    return _open(dvs_path, grb_content, indices, simulation_start_time, time_unit)


def open_conc(
    ucn_path: FilePath,
    grb_path: FilePath,
    dry_nan: bool = False,
    simulation_start_time: Optional[np.datetime64] = None,
    time_unit: Optional[str] = "d",
) -> GridDataArray:
    """
    Open Modflow6 "Unformatted Concentration" (.ucn) file.

    The data is lazily read per timestep and automatically converted into
    (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV respectively.
    The conversion is done via the information stored in the Binary Grid file
    (GRB).

    Parameters
    ----------
    ucn_path: Union[str, pathlib.Path]
    grb_path: Union[str, pathlib.Path]
    dry_nan: bool, default value: False.
        Whether to convert dry values to NaN.
    simulation_start_time : Optional datetime
        The time and date correpsonding to the beginning of the simulation.
        Use this to convert the time coordinates of the output array to
        calendar time/dates. time_unit must also be present if this argument is present.
    time_unit: Optional str
        The time unit MF6 is working in, in string representation.
        Only used if simulation_start_time was provided.
        Admissible values are:
        ns -> nanosecond
        ms -> microsecond
        s -> second
        m -> minute
        h -> hour
        d -> day
        w -> week
        Units "month" or "year" are not supported, as they do not represent unambiguous timedelta values durations.

    Returns
    -------
    concentration: Union[xr.DataArray, xu.UgridDataArray]

    Examples
    --------
    When you completed MODFLOW6 simulation, you can open the ucn file:

    >>> concentration = imod.mf6.open_conc("my-model.ucn", "my-model.grb")

    This will return a DataArray with the concentration values for each cell in
    the model, with dimensions ("time", "layer", "y", "x"). The time coordinate
    will only be the incremental time steps (1.0, 2.0, ...), unless you provide
    the ``simulation_start_time`` and ``time_unit`` arguments. This will convert
    the time coordinate to calendar time, starting from the provided simulation
    start time:

    >>> concentration = imod.mf6.open_conc(
    ...     "my-model.ucn",
    ...     "my-model.grb",
    ...     simulation_start_time=np.datetime64("2023-01-01T00:00:00"),
    ...     time_unit="d",
    ... )

    """
    grb_content = read_grb(grb_path)
    grb_content["name"] = "concentration"
    distype = grb_content["distype"]
    _open = _get_function(_OPEN_HDS, distype)
    return _open(ucn_path, grb_content, dry_nan, simulation_start_time, time_unit)


def open_hds_like(
    path: FilePath,
    like: Union[xr.DataArray, xu.UgridDataArray],
    dry_nan: bool = False,
) -> GridDataArray:
    """
    Open modflow6 heads (.hds) file.

    The data is lazily read per timestep and automatically converted into
    DataArrays. Shape and coordinates are inferred from ``like``.

    Parameters
    ----------
    hds_path: Union[str, pathlib.Path]
    like: Union[xr.DataArray, xu.UgridDataArray]
    dry_nan: bool, default value: False.
        Whether to convert dry values to NaN.

    Returns
    -------
    head: Union[xr.DataArray, xu.UgridDataArray]
    """
    # TODO: check shape with hds metadata.
    if isinstance(like, xr.DataArray):
        d = dis.grid_info(like)
        return dis.open_hds(path, d, dry_nan)

    elif isinstance(like, xu.UgridDataArray):
        d = disv.grid_info(like)
        return disv.open_hds(path, d, dry_nan)

    else:
        raise TypeError(
            "like should be a DataArray or UgridDataArray, "
            f"received instead {type(like)}"
        )


def open_cbc(
    cbc_path: FilePath,
    grb_path: FilePath,
    flowja: bool = False,
    simulation_start_time: Optional[np.datetime64] = None,
    time_unit: Optional[str] = "d",
    merge_to_dataset: bool = False,
) -> GridDataset | Dict[str, GridDataArray]:
    """
    Open modflow6 cell-by-cell (.cbc) file.

    The data is lazily read per timestep and automatically converted into
    (dense) xr.DataArrays or xu.UgridDataArrays, for DIS and DISV respectively.
    The conversion is done via the information stored in the Binary Grid file
    (GRB).

    The ``flowja`` argument controls whether the flow-ja-face array (if present)
    is returned in grid form as "as is". By default ``flowja=False`` and the
    array is returned in "grid form", meaning:

        * DIS: in right, front, and lower face flow. All flows are placed in
          the cell.
        * DISV: in horizontal and lower face flow.the horizontal flows are
          placed on the edges and the lower face flow is placed on the faces.

    When ``flowja=True``, the flow-ja-face array is returned as it is found in
    the CBC file, with a flow for every cell to cell connection. Additionally,
    a ``connectivity`` DataArray is returned describing for every cell (n) its
    connected cells (m).

    Parameters
    ----------
    cbc_path: str, pathlib.Path
        Path to the cell-by-cell flows file
    grb_path: str, pathlib.Path
        Path to the binary grid file
    flowja: bool, default value: False
        Whether to return the flow-ja-face values "as is" (``True``) or in a
        grid form (``False``).
    simulation_start_time : Optional datetime
        The time and date correpsonding to the beginning of the simulation.
        Use this to convert the time coordinates of the output array to
        calendar time/dates. time_unit must also be present if this argument is present.
    time_unit: Optional str
        The time unit MF6 is working in, in string representation.
        Only used if simulation_start_time was provided.
        Admissible values are:
        ns -> nanosecond
        ms -> microsecond
        s -> second
        m -> minute
        h -> hour
        d -> day
        w -> week
        Units "month" or "year" are not supported, as they do not represent unambiguous timedelta values durations.
    merge_to_dataset: bool, default value: False
        Merge output to dataset.

    Returns
    -------
    cbc_content: xr.Dataset | Dict[str, xr.DataArray]
        DataArray contains float64 data of the budgets, with dimensions ("time",
        "layer", "y", "x").

    Examples
    --------

    Open a cbc file:

    >>> import imod
    >>> cbc_content = imod.mf6.open_cbc("budgets.cbc", "my-model.grb")

    Check the contents:

    >>> print(cbc_content.keys())

    Get the drainage budget, compute a time mean for the first layer:

    >>> drn_budget = cbc_content["drn]
    >>> mean = drn_budget.sel(layer=1).mean("time")

    This will return a DataArray with the drain budget values for each cell in the
    model, with dimensions ("time", "layer", "y", "x"). The time coordinate will
    only be the incremental time steps (1.0, 2.0, ...), unless you provide the
    ``simulation_start_time`` and ``time_unit`` arguments. This will convert the
    time coordinate to calendar time, starting from the provided simulation
    start time:

    >>> cbc_content = imod.mf6.open_cbc(
    ...     "my-model.cbc",
    ...     "my-model.grb",
    ...     simulation_start_time=np.datetime64("2023-01-01T00:00:00"),
    ...     time_unit="d",
    ... )
    """
    grb_content = read_grb(grb_path)
    distype = grb_content["distype"]
    _open = _get_function(_OPEN_CBC, distype)
    cbc = _open(cbc_path, grb_content, flowja, simulation_start_time, time_unit)
    if merge_to_dataset:
        return merge_with_dictionary(cbc)
    return cbc