# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['raster_to_xarray', 'read_images_into_stacked_array', 'mapcycle_to_xarray',
           'read_mapcycles_into_stacked_array']

# Internal Cell
from pathlib import Path

import xarray as xr

# Cell
def raster_to_xarray(fpath, chunk_scale=2):
    """Read raster image into an xarray.DataArray.

    Using the `chunks` keyword in the open_rasterio method
    activates the return of an out-of-memory virtual array instead
    of the in-memory xarray.DataArray

    fpath: pathlib.Path, str
    chunk_scale: int
        Multiplier for the hardcoded 2048/1024 chunk-sizes for
        x/y axes.
    """
    fpath = Path(fpath)
    return xr.open_rasterio(
        fpath, chunks={"x": chunk_scale * 2048, "y": chunk_scale * 1024}
    )

# Cell
def read_images_into_stacked_array(image_paths, name, chunk_scale=2):
    arrays = [raster_to_xarray(p, chunk_scale) for p in image_paths]
    stack = xr.concat(arrays, "time")
    stack.name = name
    return stack

# Cell
def mapcycle_to_xarray(fpath, chunk_scale=2):
    """Read an Diviner map cycle rasterio image into a dask.array.

    The cycle identifier will be read from the filename and added
    to the xarray as a coordinate value.

    fpath: str, pathlib.Path
    chunk_scale: int
        Scaling the chunk
    """
    fpath = Path(fpath)
    da = raster_to_xarray(fpath)
    cycle = int(fpath.name.split("_")[4][:-1])
    da = da.assign_coords(band=[cycle])
    da = da.rename({"band": "mapcycle", "y": "lat", "x": "lon"})
    da.lat.attrs["long_name"] = "Latitude"
    da.lon.attrs["long_name"] = "Longitude"
    for coord in ["lat", "lon"]:
        da[coord].attrs["units"] = "deg"
    return da

# Cell
def read_mapcycles_into_stacked_array(image_paths, name, chunk_scale=2):
    with xr.set_options(keep_attrs=True):
        arrays = [mapcycle_to_xarray(p, chunk_scale) for p in image_paths]
        stack = xr.concat(arrays, "mapcycle")
        stack.name = name
        return stack