"""General utility functions"""

import numpy as np
from skimage import color

def crop_zeros(x, y=None, value=0.0):
    """Crop an array, `x` to its non-zero elements, or optionally to those of a second array `y`
    
    Parameters
    ----------
    x : ndarray
        The array that should be cropped.
    y: ndarray, optional
        The array (identical in shape to `x`) by whose non-zero elements `x` should be cropped. If `None`, will use non-zero elements of `x`.
    value: float
        The (background) value that should be cropped out.
    
    Returns
    -------
    ndarray
        `x`, cropped to the non-zero elements of `x` or, if specified, `y`
    
    Examples
    --------
    >>> a = np.arange(16).reshape((4, 4))
    >>> a[0, :] = 0
    >>> a[:, 3] = 0
    >>> crop_zeros(a)
    array([[ 4,  5,  6],
           [ 8,  9, 10],
           [12, 13, 14]])
    >>> b = np.arange(16).reshape(4, 4)
    >>> crop_zeros(b, a)
    array([[ 4,  5,  6],
           [ 8,  9, 10],
           [12, 13, 14]])
    """
    # find non-zero elements
    nonzero_idx = np.where(x!=value) if y is None else np.where(y!=value)
    # create list of slices by which x should be indexed
    crop_idx = [slice(np.min(d), np.max(d)+1) for d in nonzero_idx]
    x_crop = x[tuple(crop_idx)]
    return(x_crop)

def pad_translate_mat(x, dim0, dim1, value=0):
    """Translate a matrix with padding
    
    Parameters
    ----------
    x : ndarray
        The matrix to pad.
    dim0 : float
        The translation to apply in axis 0.
    dim1 : float
        The translation to apply in axis 1.
    value : float
        The value to pad with when inserting rows/columns.
    
    Returns
    -------
    ndarray
        `x`, with values translated.
    
    Examples
    --------
    >>> a = np.zeros((4, 4))
    >>> a[1, :] = 1
    >>> a
    array([[0., 0., 0., 0.],
           [0., 0., 0., 0.],
           [1., 1., 1., 1.],
           [0., 0., 0., 0.]])
    >>> pad_translate_mat(a, 1, 0)
    array([[0., 0., 0., 0.],
           [1., 1., 1., 1.],
           [0., 0., 0., 0.],
           [0., 0., 0., 0.]])
    """
    # if dim0 < 0:
    #     x = np.pad(x, [(abs(dim0), value), (value, value)])
    #     x = x[0:dim0, :]
    # elif dim0 > 0:
    #     x = np.pad(x, [(value, dim0), (value, value)])
    #     x = x[dim0:, :]
    # if dim1 < 0:
    #     x = np.pad(x, [(value, value), (abs(dim1), value)])
    #     x = x[:, 0:dim1]
    # elif dim1 > 0:
    #     x = np.pad(x, [(value, value), (value, dim1)])
    #     x = x[:, dim1:]

    other_dims = [(0,0)] * (len(x.shape) - 2)

    if dim0 < 0:
        pad0 = [(abs(dim0), value), (value, value)]
        pad0.extend(other_dims)
        x = np.pad(x, pad0)
        x = x[0:dim0, :]
    elif dim0 > 0:
        pad0 = [(value, dim0), (value, value)]
        pad0.extend(other_dims)
        x = np.pad(x, pad0)
        x = x[dim0:, :]
    if dim1 < 0:
        pad1 = [(value, value), (abs(dim1), value)]
        pad1.extend(other_dims)
        x = np.pad(x, pad1)
        x = x[:, 0:dim1]
    elif dim1 > 0:
        pad1 = [(value, value), (value, dim1)]
        pad1.extend(other_dims)
        x = np.pad(x, pad1)
        x = x[:, dim1:]

    return(x)

def pad_for_translation(mat1, mat2, pad=True, constant_values=0):
    """Get two matrices to be the same size, with zero-padding for translation. Space for translation will mean that both matrices are made to be 9 times as large (3x in both dimensions) as the largest of the two.
    
    Parameters
    ----------
    mat1 : ndarray
    mat2 : ndarray
    pad : bool
        If `True`, will pad with space for translation. If `False`, will just pad to be the same size with the smaller matrix's non-zero elements centred on those of the larger.
    constant_values : float or list or array
        Passed to `np.pad`. Will usually want to be the value for the background (default = 0).
    
    Returns
    -------
    tuple
        A tuple in the form `(mat1, mat2)`, containing the padded matrices.

    Examples
    --------
    >>> a = np.arange(16).reshape((4, 4))
    >>> b = np.arange(9).reshape((3, 3))
    >>> a_pad, b_pad = pad_for_translation(a, b)
    >>> (a_pad.shape, b_pad.shape)
    """
    assert len(mat1.shape)==2
    assert len(mat2.shape)==2
    # get the max dimensions of both
    max_dims = [np.max([mat1.shape[i], mat2.shape[i]]) for i in range(2)]
    # get the amount of padding needed
    mat1_diff = np.abs([mat1.shape[i] - max_dims[i] for i in range(len(max_dims))])
    mat2_diff = np.abs([mat2.shape[i] - max_dims[i] for i in range(len(max_dims))])
    # if any odd numbers, add 1 to the diff for each array    
    mat1_modulo = mat1_diff % 2
    mat2_modulo = mat2_diff % 2
    mat1_diff += mat1_modulo
    mat2_diff += mat2_modulo
    # account for possible difference in size
    mat1 = np.pad(mat1, [(0, mat2_modulo[0]), (0, mat2_modulo[1])], constant_values=constant_values)
    mat2 = np.pad(mat2, [(0, mat1_modulo[0]), (0, mat1_modulo[1])], constant_values=constant_values)
    # pad to be the same size
    mat1_pad_size = [(int(mat1_diff[i]/2), int(mat1_diff[i]/2)) for i in range(len(mat1_diff))]
    mat2_pad_size = [(int(mat2_diff[i]/2), int(mat2_diff[i]/2)) for i in range(len(mat2_diff))]
    mat1 = np.pad(mat1, mat1_pad_size, constant_values=constant_values)
    mat2 = np.pad(mat2, mat2_pad_size, constant_values=constant_values)
    # check they are the same size
    assert mat1.shape == mat2.shape

    if pad:
        # now pad to allow all possible translations
        mat1 = np.pad(mat1, [(mat1.shape[0], mat1.shape[0]), (mat1.shape[1], mat1.shape[1])], constant_values=constant_values)
        mat2 = np.pad(mat2, [(mat2.shape[0], mat2.shape[0]), (mat2.shape[1], mat2.shape[1])], constant_values=constant_values)
        # and pad the starts to make even
        mat1_modulo_b = np.array(mat1.shape) % 2
        mat2_modulo_b = np.array(mat2.shape) % 2
        mat1 = np.pad(mat1, [(mat1_modulo_b[0], 0), (mat1_modulo_b[1], 0)], constant_values=constant_values)
        mat2 = np.pad(mat2, [(mat2_modulo_b[0], 0), (mat2_modulo_b[1], 0)], constant_values=constant_values)
    # return the two matrices as a tuple
    return(mat1, mat2)

def chunk_list(l, n):
    """Chunk a list, `l`, into `n` chunks of (as close to as possible) equal length, keeping original sequential order. Written by https://stackoverflow.com/a/54802737
    
    Parameters
    ----------
    l : list
        The list to be chunked.
    n : int
        The number of chunks to produce.

    Returns
    -------
    generator object
        A generator object with the list chunked.
    """
    d, r = divmod(len(l), n)

    for i in range(n):
        si = (d+1)*(i if i < r else r) + d*(0 if i < r else i - r)
        yield l[si:si+(d+1 if i < r else d)]

def rotate_rgb_hue(rgb, hue_shift=0.5):
    """Rotate the hue of the colour channel of a numpy array. RGB should be stored in the final dimension.

    Parameters
    ----------
    rgb : np.array
        The array with colour values whose hue should be shifted.
    hue_shift : float
        A value from 0 to 1 that dictates the degrees to which the hue should be rotated. Degrees are equivalent to `hue_shuft*360`.

    Returns
    -------
    np.array
        The original np.array, but with colour values in the rgb channel rotated.
    """
    hsv = color.rgb2hsv(rgb)
    hsv[:, :, 0] += hue_shift
    hsv[:, :, 0][hsv[:, :, 0] > 1] -= hsv[:, :, 0].min()
    return color.hsv2rgb(hsv)

def logistic(x):
    x = np.array(x)
    return np.log( (x+1) / (1-x) )

def inv_logistic(x):
    x = np.array(x)
    return ( np.exp(x)-1 ) / ( np.exp(x)+1 )