"""Various useful functions."""
import numpy as np
import matplotlib.pyplot as plt


def read_data():
    """Reads the exercise data as numpy arrays and scalars

    Returns
    -------
    a tuple of (gridsize_m, elevation, mask)

    The variable “elevation” contains a Digital Elevation Model (DEM, unit: m)
    of a region around a glacier.
    The variable “mask” contains a mask identifying the location of the glacier.
    The variable “dx” indicates the grid spacing (the resolution)
    of the DEM (unit: m).
    """
    import xarray as xr
    ds = xr.open_dataset('glacier.nc')
    return ds.gridsize_m, ds.dem.values, ds.glacier_mask.values


def plot_numpy_array_with_colorbar(a):
    """Plots a numpy array with a colorbar for quick analysis"""
    plt.imshow(a)
    plt.colorbar()
    plt.show()


def find_max_indices_in_2d_array(data):
    """Finds the location of the maximum value in a 2D array.

    Parameters
    ----------
    data : ndarray of dimensions 2
        the data to find the max from

    Returns
    -------
    (j, i) : tuple of i, j indices in the array where the max is located
    """
    if len(data.shape) != 2:
        raise ValueError('arrays of shape different than 2 is not supported yet')
    return np.unravel_index(data.argmax(), data.shape)


def zoom_into_2d_array(data, j, i, half_size=20):
    """Zoom into a 2D array at a given location.

    Parameters
    ----------
    data : ndarray of dimensions 2
        the data to zoom into
    j : int
        the y-axis indice to zoom into
    i : int
        the x-axis indice to zoom into
    half_size : int
        the half-size of the zoom window

    Returns
    -------
    a quadratic 2D ndarray of size (2*half_size+1)
    """
    if len(data.shape) != 2:
        raise ValueError('arrays of shape different than 2 is not supported yet')
    sy, sx = data.shape
    if (j - half_size < 0) or (i - half_size < 0) or (j + half_size >= sy) or (i + half_size >= sx):
        raise ValueError('half_size does not match size of array')
    return data[j - half_size:j + half_size + 1, i - half_size:i + half_size + 1]


if __name__ == '__main__':

    # Quick and dirty tests
    # TODO: will make formal tests later

    dx, dem, mask = read_data()

    plot_numpy_array_with_colorbar(mask)
    
    plot_numpy_array_with_colorbar(dem)

    j, i = find_max_indices_in_2d_array(dem)
    plot_numpy_array_with_colorbar(zoom_into_2d_array(dem, j, i))