_io.py

"""Functions for loading and saving cubes."""

from __future__ import annotations

import copy
import logging
import os
import warnings
from itertools import groupby
from pathlib import Path
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

import fsspec
import iris
import ncdata
import xarray as xr
import yaml
from iris.cube import Cube, CubeList

from esmvalcore._task import write_ncl_settings
from esmvalcore.exceptions import ESMValCoreLoadWarning
from esmvalcore.io.local import LocalFile
from esmvalcore.io.protocol import DataElement
from esmvalcore.iris_helpers import dataset_to_iris

if TYPE_CHECKING:
    from collections.abc import Sequence

    from dask.delayed import Delayed

logger = logging.getLogger(__name__)

GLOBAL_FILL_VALUE = 1e20

DATASET_KEYS = {
    "mip",
}
VARIABLE_KEYS = {
    "reference_dataset",
    "alternative_dataset",
}


def _drop_range_attributes(cube: Cube) -> Cube:
    """Drop range related attributes from cube and its components."""
    drop_attrs = (
        "actual_range",
        "valid_max",
        "valid_min",
        "valid_range",
    )
    cube = cube.copy()
    for attr in drop_attrs:
        cube.attributes.pop(attr, None)
        for coord in cube.dim_coords:
            coord.attributes.pop(attr, None)
        for aux_coord in cube.aux_coords:
            aux_coord.attributes.pop(attr, None)
        for cell_measure in cube.cell_measures():
            cell_measure.attributes.pop(attr, None)
        for ancillary_variable in cube.ancillary_variables():
            ancillary_variable.attributes.pop(attr, None)
    return cube


def load(
    file: str
    | Path
    | DataElement
    | Cube
    | CubeList
    | xr.Dataset
    | ncdata.NcData,
    ignore_warnings: list[dict[str, Any]] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
) -> CubeList:
    """Load Iris cubes.

    Parameters
    ----------
    file:
        File to be loaded. If ``file`` is already a loaded dataset, return it
        as a :class:`~iris.cube.CubeList`.
        File as ``Path`` object could be a Zarr store.
    ignore_warnings:
        Keyword arguments passed to :func:`warnings.filterwarnings` used to
        ignore warnings issued by :func:`iris.load_raw`. Each list element
        corresponds to one call to :func:`warnings.filterwarnings`.
    backend_kwargs:
        Dict to hold info needed by storage backend e.g. to access
        a PRIVATE S3 bucket containing object stores (e.g. netCDF4 files);
        needed by ``fsspec`` and its extensions e.g. ``s3fs``, so
        most of the times this will include ``storage_options``. Note that Zarr
        files are opened via ``http`` extension of ``fsspec``, so no need
        for ``storage_options`` in that case (ie anon/anon). Currently only used
        in Zarr file opening.

    Returns
    -------
    iris.cube.CubeList
        Loaded cubes.

    Raises
    ------
    ValueError
        Cubes are empty.
    TypeError
        Invalid type for ``file``.

    """
    if isinstance(file, DataElement):
        cubes = file.to_iris()
    elif isinstance(file, (str, Path)):
        extension = (
            file.suffix
            if isinstance(file, Path)
            else os.path.splitext(file)[1]
        )
        if "zarr" not in extension:
            local_file = LocalFile(file)
            local_file.ignore_warnings = ignore_warnings
            cubes = local_file.to_iris()
        else:
            cubes = _load_zarr(
                file,
                ignore_warnings=ignore_warnings,
                backend_kwargs=backend_kwargs,
            )
    elif isinstance(file, Cube):
        cubes = CubeList([file])
    elif isinstance(file, CubeList):
        cubes = file
    elif isinstance(file, (xr.Dataset, ncdata.NcData)):
        cubes = dataset_to_iris(file, ignore_warnings=ignore_warnings)
    else:
        msg = (
            f"Expected type str, pathlib.Path, iris.cube.Cube, "
            f"iris.cube.CubeList, xarray.Dataset, or ncdata.NcData for file, "
            f"got type {type(file)}"
        )
        raise TypeError(msg)

    if not cubes:
        msg = f"{file} does not contain any data"
        raise ValueError(msg)

    for cube in cubes:
        if "source_file" not in cube.attributes:
            warn_msg = (
                f"Cube {cube.summary(shorten=True)} loaded from\n{file}\ndoes "
                f"not contain attribute 'source_file' that points to original "
                f"file path, please make sure to add it prior to loading "
                f"(preferably during the preprocessing step 'fix_file')"
            )
            warnings.warn(warn_msg, ESMValCoreLoadWarning, stacklevel=2)

    # Drop range related attributes as these are likely to be
    # invalidated by preprocessing the data.
    return CubeList(_drop_range_attributes(cube) for cube in cubes)


def _load_zarr(
    file: str | Path,
    ignore_warnings: list[dict[str, Any]] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
) -> CubeList:
    # note on ``chunks`` kwarg to ``xr.open_dataset()``
    # docs.xarray.dev/en/stable/generated/xarray.open_dataset.html
    # this is very important because with ``chunks=None`` (default)
    # data will be realized as Numpy arrays and transferred in memory;
    # ``chunks={}`` loads the data with dask using the engine preferred
    # chunk size, generally identical to the formats chunk size. If not
    # available, a single chunk for all arrays; testing shows this is the
    # "best guess" compromise for typically CMIP-like chunked data.
    # see https://github.com/pydata/xarray/issues/10612 and
    # https://github.com/pp-mo/ncdata/issues/139

    time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
    open_kwargs = {
        "consolidated": False,
        "decode_times": time_coder,
        "engine": "zarr",
        "chunks": {},
        "backend_kwargs": backend_kwargs,
    }

    # Case 1: Zarr store is on remote object store
    # file's URI will always be either http or https
    if urlparse(str(file)).scheme in ["http", "https"]:
        # basic test that opens the Zarr/.zmetadata file for Zarr2
        # or Zarr/zarr.json for Zarr3
        fs = fsspec.filesystem("http")
        valid_zarr = True
        try:
            fs.open(str(file) + "/zarr.json", "rb")  # Zarr3
        except Exception:  # noqa: BLE001
            try:
                fs.open(str(file) + "/.zmetadata", "rb")  # Zarr2
            except Exception:  # noqa: BLE001
                valid_zarr = False
        # we don't want to catch any specific aiohttp/fsspec exception
        # bottom line is that that file has issues, so raise
        if not valid_zarr:
            msg = (
                f"File '{file}' can not be opened as Zarr file at the moment."
            )
            raise ValueError(msg)

        open_kwargs["consolidated"] = True
        zarr_xr = xr.open_dataset(file, **open_kwargs)
    # Case 2: Zarr store is local to the file system
    else:
        zarr_xr = xr.open_dataset(file, **open_kwargs)

    # avoid possible
    # ValueError: Object has inconsistent chunks along dimension time.
    # This can be fixed by calling unify_chunks().
    # when trying to access the ``chunks`` store
    zarr_xr = zarr_xr.unify_chunks()

    return dataset_to_iris(zarr_xr, ignore_warnings=ignore_warnings)


def save(  # noqa: C901
    cubes: Sequence[Cube],
    filename: Path | str,
    optimize_access: str = "",
    compress: bool = False,
    alias: str = "",
    compute: bool = True,
    **kwargs: Any,
) -> Delayed | None:
    """Save iris cubes to file.

    Parameters
    ----------
    cubes:
        Data cubes to be saved

    filename:
        Name of target file

    optimize_access:
        Set internal NetCDF chunking to favour a reading scheme

        Values can be map or timeseries, which improve performance when
        reading the file one map or time series at a time.
        Users can also provide a coordinate or a list of coordinates. In that
        case the better performance will be avhieved by loading all the values
        in that coordinate at a time

    compress:
        Use NetCDF internal compression.

    alias:
        Var name to use when saving instead of the one in the cube.

    compute : bool, default=True
        Default is ``True``, meaning complete the file immediately, and return
        ``None``.

        When ``False``, create the output file but don't write any lazy array
        content to its variables, such as lazy cube data or aux-coord points
        and bounds.  Instead return a :class:`dask.delayed.Delayed` which, when
        computed, will stream all the lazy content via :meth:`dask.store`, to
        complete the file.  Several such data saves can be performed in
        parallel, by passing a list of them into a :func:`dask.compute` call.

    **kwargs:
        See :func:`iris.fileformats.netcdf.saver.save` for additional
        keyword arguments.

    Returns
    -------
    :class:`dask.delayed.Delayed` or :obj:`None`
        A delayed object that can be used to save the data in the cube.

    Raises
    ------
    ValueError
        cubes is empty.
    """
    if not cubes:
        msg = f"Cannot save empty cubes '{cubes}'"
        raise ValueError(msg)

    if Path(filename).suffix.lower() == ".nc":
        kwargs["compute"] = compute

    # Rename some arguments
    kwargs["target"] = filename
    kwargs["zlib"] = compress

    dirname = os.path.dirname(filename)
    if not os.path.exists(dirname):
        os.makedirs(dirname)

    if os.path.exists(filename) and all(
        cube.has_lazy_data() for cube in cubes
    ):
        logger.debug(
            "Not saving cubes %s to %s to avoid data loss. "
            "The cube is probably unchanged.",
            cubes,
            filename,
        )
        return None

    for cube in cubes:
        logger.debug(
            "Saving cube:\n%s\nwith %s data to %s",
            cube,
            "lazy" if cube.has_lazy_data() else "realized",
            filename,
        )
    if optimize_access:
        cube = cubes[0]
        if optimize_access == "map":
            dims = set(
                cube.coord_dims("latitude") + cube.coord_dims("longitude"),
            )
        elif optimize_access == "timeseries":
            dims = set(cube.coord_dims("time"))
        else:
            dims = {
                dim
                for coord_name in optimize_access.split(" ")
                for dim in cube.coord_dims(coord_name)
            }

        kwargs["chunksizes"] = tuple(
            length if index in dims else 1
            for index, length in enumerate(cube.shape)
        )

    kwargs["fill_value"] = GLOBAL_FILL_VALUE
    if alias:
        for cube in cubes:
            logger.debug(
                "Changing var_name from %s to %s",
                cube.var_name,
                alias,
            )
            cube.var_name = alias

    # Ignore some warnings when saving
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message=(
                ".* is being added as CF data variable attribute, but .* "
                "should only be a CF global attribute"
            ),
            category=UserWarning,
            module="iris",
        )
        return iris.save(cubes, **kwargs)


def _get_debug_filename(filename, step):
    """Get a filename for debugging the preprocessor."""
    dirname = os.path.splitext(filename)[0]
    if os.path.exists(dirname) and os.listdir(dirname):
        num = int(sorted(os.listdir(dirname)).pop()[:2]) + 1
    else:
        num = 0
    return os.path.join(dirname, f"{num:02}_{step}.nc")


def _sort_products(products):
    """Sort preprocessor output files by their order in the recipe."""
    return sorted(
        products,
        key=lambda p: (
            p.attributes.get("recipe_dataset_index", 1e6),
            p.attributes.get("dataset", ""),
        ),
    )


def write_metadata(products, write_ncl=False):
    """Write product metadata to file."""
    output_files = []
    for output_dir, prods in groupby(
        products,
        lambda p: os.path.dirname(p.filename),
    ):
        sorted_products = _sort_products(prods)
        metadata = {}
        for product in sorted_products:
            if isinstance(product.attributes.get("exp"), (list, tuple)):
                product.attributes = dict(product.attributes)
                product.attributes["exp"] = "-".join(product.attributes["exp"])
            if "original_short_name" in product.attributes:
                del product.attributes["original_short_name"]
            metadata[product.filename] = product.attributes

        output_filename = os.path.join(output_dir, "metadata.yml")
        output_files.append(output_filename)
        with open(output_filename, "w", encoding="utf-8") as file:
            yaml.safe_dump(metadata, file)
        if write_ncl:
            output_files.append(_write_ncl_metadata(output_dir, metadata))

    return output_files


def _write_ncl_metadata(output_dir, metadata):
    """Write NCL metadata files to output_dir."""
    variables = [copy.deepcopy(v) for v in metadata.values()]

    info = {"input_file_info": variables}

    # Split input_file_info into dataset and variable properties
    # dataset keys and keys with non-identical values will be stored
    # in dataset_info, the rest in variable_info
    variable_info = {}
    info["variable_info"] = [variable_info]
    info["dataset_info"] = []
    for variable in variables:
        dataset_info = {}
        info["dataset_info"].append(dataset_info)
        for key in variable:
            dataset_specific = any(
                variable[key] != var.get(key, object()) for var in variables
            )
            if (
                dataset_specific or key in DATASET_KEYS
            ) and key not in VARIABLE_KEYS:
                dataset_info[key] = variable[key]
            else:
                variable_info[key] = variable[key]

    filename = os.path.join(
        output_dir,
        variable_info["short_name"] + "_info.ncl",
    )
    write_ncl_settings(info, filename)

    return filename