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

__all__ = ['memmap_binary', 'read_hpar_binary', 'basepath', 'lons_path', 'lats_path', 'hpar_path', 'HReader']

# Cell
from pathlib import Path

import dask.array as da
import dask
import hvplot.xarray
import numpy as np
import xarray as xr

basepath = Path("/luna4/maye/h_parameter_maps")
lons_path = basepath / "hpar_global128ppd_v1c.lon.bin"
lats_path = basepath / "hpar_global128ppd_v1c.lat.bin"
hpar_path = basepath / "hpar_global128ppd_v1c.bin"


def memmap_binary(fname):
    # one extra element for both lats (the 0-line) and lons (360==0)
    cols = 128 * 360 + 1
    rows = 128 * 140 + 1
    # this doesn't read, just maps the diskfile into a virtual array
    return np.memmap(fname, np.float32, mode="r", shape=(rows, cols), order="F")


def read_hpar_binary(fname=hpar_path):
    """Reading Paul Hayne's binary global 128 ppd H-parameter map.

    Parameters
    ----------
    fname: str, pathlib.Path
        Path to binary 128 ppd map file binary (ending in .bin)
    """
    if Path(fname).suffix != ".bin":
        print("not an expected filename ending in .bin .")
        return

    arr = memmap_binary(hpar_path)

    # create a dask array
    a = da.from_array(arr)

    # replacing the NAN values with actual np.nan
    nan_value = -32768
    a[a == nan_value] = np.nan

    # making up the lon/lat grid. Might need adaptation to match other global data
    lons = np.linspace(0, 360, arr.shape[1])
    lats = np.linspace(70, -70, arr.shape[0])

    # create xarray DataArray from the dask array (thankfully supported, but not documented)
    img = xr.DataArray(a, dims=["lat", "lon"], coords={"lat": lats, "lon": lons})
    img.name = "H"
    return img

# Cell
class HReader:
    """Data reader class for H parameter map.

    It accesses preproduced data for 128 ppd for DEM, slope, and aspect,
    located on the luna4 disk.

    It uses virtual dask.Arrays so that virtually no memory is consumed until you
    resolve a chain of operations with the `.compute()` call.

    Attributes
    ----------
    H: xarray.DataArray
        H parameter map
    """

    def __init__(self, lat_limit=None):
        self.lat_limit = lat_limit

        arr = memmap_binary(hpar_path)
        # dask array:
        a = da.from_array(arr)
        # replace NAN values with np.nan
        a[a == -32768] = np.nan

        # As per Paul's README:
        lats = np.linspace(70, -70, arr.shape[0], endpoint=True)
        lons = np.linspace(0, 360, arr.shape[1], endpoint=True)

        # create xarray DataArray from the dask array (thankfully supported, but not documented)
        img = xr.DataArray(a, dims=["lat", "lon"], coords={"lat": lats, "lon": lons})
        img.name = "H"
        img.attrs["long_name"] = "H Parameter"
        img.attrs["units"] = "H units :-P"
        self.img = img

    def slice_lat(self, lat):
        """Return the map `data` constrained to lat <= `lat`.

        Parameters
        ----------
        data: {'dem', 'slope','aspect'}
            String that choses which data product should be constrained.
        lat: int, float
            Limiting latitude value.
        """
        s = slice(lat, -lat)
        return self.img.sel(lat=s, drop=True)

    def convert_to_lon180(self):
        "Switch image to -180/180 longitude system."
        p = self.img
        p.coords["lon"] = ((p.lon + 180) % 360) - 180
        with dask.config.set(**{"array.slicing.split_large_chunks": False}):
            self.img = p.sortby(p.lon)

    def convert_to_lon360(self):
        "Switch image to 360 longitude system."
        p = self.img
        p.coords["lon"] = p.coords["lon"] % 360
        with dask.config.set(**{"array.slicing.split_large_chunks": False}):
            self.img = p.sortby(p.lon)

    def assign_new_latlon(self, lat, lon):
        "make sure longitude layout matches!"
        p = self.img
        p.coords["lon"] = lon
        p.coords["lat"] = lat
        self.img = p.sortby(p.lon)

    def get_H_by_pixel(self, ilat, ilon):
        "Note: Decided to not apply offset!"
        val = self.img.isel(lat=ilat, lon=ilon)
        return float(val)

    def get_H_by_coord(self, lat, lon):
        "Note: Decided to not apply offset!"
        val = self.img.sel(lat=lat, lon=lon, method="nearest")
        return float(val)

    def plot_dem(self, lat_min, lon_min, dlat=1, dlon=1, lat_max=None, lon_max=None):
        sliced = self.get_slice("dem", lat_min, lon_min, dlat, dlon, lat_max, lon_max)
        return sliced.hvplot(cmap="viridis", aspect="equal", title="DEM")

    def plot_slope(self, lat_min, lon_min, dlat=1, dlon=1, lat_max=None, lon_max=None):
        sliced = self.get_slice("slope", lat_min, lon_min, dlat, dlon, lat_max, lon_max)
        return sliced.hvplot(cmap="inferno", aspect="equal", title="Slope [degrees]")

    def plot_aspect(self, lat_min, lon_min, dlat=1, dlon=1, lat_max=None, lon_max=None):
        sliced = self.get_slice(
            "aspect", lat_min, lon_min, dlat, dlon, lat_max, lon_max
        )
        return sliced.hvplot(cmap="twilight", aspect="equal", title="Azimuth [degrees]")

    def get_slice(
        self, obj, lat_min, lon_min, dlat=None, dlon=None, lat_max=None, lon_max=None
    ):
        """Slice a rectangular data tile out of the map.

        Parameters
        ----------
        obj: {'dem','slope','aspect'}
            Determines the object that will be sliced
        lat_min: float
            Lower left corner latitude of slice.
        lon_min: float
            Lower left corner longitude of slice.
        dlat: float
            Delta to be added to `lat_min`
        dlon: float
            Delta lon to be added to `lon_min`
        lat_max: float
            Alternative for dlat, to set upper end of latitude interval.
        lon_max: float
            Alternative for dlon, to set right end of longitude interval.

        Returns
        -------
        xarray.DataArray
            Sliced for the provided coordinate values.
        """
        data = getattr(self, obj)
        if lon_max is None:
            lon_max = lon_min + dlon
        if lat_max is None:
            lat_max = lat_min + dlat
        return data.sel(lat=slice(lat_max, lat_min), lon=slice(lon_min, lon_max))