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

__all__ = ['interpolate_screen_pos', 'interpolate_checker_pos', 'split_eye_events', 'get_spherical_map',
           'apply_spherical_map']

# Cell
import numpy as np
import scipy.interpolate as interpolate
from sklearn import cluster

# Cell
def interpolate_screen_pos(screen_pos, xnew, ynew, kind='linear'):
    """
    Interpolate the position of the xnew and ynew pixels from the original screen_pos.
    `interpolate_checker_pos` should be used instead as it's more user friendly.
    params:
        - screen_pos: Screen positions in shape (17, 10, 2) obtained from calibration (from 80x80 pixel checker corners on a 1280x720px screen)
        - xnew: New pixels indexes in x, in interval [0, 17[
        - ynew: New pixels indexes in y, in interval [0, 10[
    return:
        - Screen positions in shape (len(xnew),len(ynew), 2)
    """
    f = interpolate.interp2d(np.arange(17), np.arange(10), screen_pos[:,:,0].T, kind=kind)
    znew_x = f(xnew, ynew)

    f = interpolate.interp2d(np.arange(17), np.arange(10), screen_pos[:,:,1].T, kind=kind)
    znew_y = f(xnew, ynew)

    return np.stack((znew_x, znew_y), axis=-1)

def interpolate_checker_pos(screen_pos, width_box, height_box, kind='linear'):
    """
    Interpolate the centre of the checkerboard pixels from the screen calibrated position.
    params:
        - screen_pos: Screen positions in shape (17, 10, 2) obtained from calibration (from 80x80 pixel checker corners on a 1280x720px screen)
        - width_box: Width in pixel of a box
        - height_box: Height in pixel of a box
        - kind: kind of interpolation in {'linear', 'cubic', 'quintic'}
    """
    assert 1280%width_box==0, "unpredictable behaviour if 1280 is not a multiple of width_box"
    assert 720%height_box==0, "unpredictable behaviour if 720 is not a multiple of height_box"
    n_x = 1280/width_box
    n_y = 720/height_box
    xnew = np.arange(16/n_x/2, 16+16/n_x/2, 16/n_x)
    ynew = np.arange(9/n_y/2, 9+9/n_y/2, 9/n_y)
    return interpolate_screen_pos(screen_pos, xnew, ynew, kind=kind)

# Cell
def split_eye_events(eye_tracking, eps=2):
    """
    Split the record where the eye moves. Detection done with clustering on X,Y and time of the eye position.

    params:
        - eye_tracking: Eye traking array of the ellipse fit, in shape (t, (x,y,width,height,angle))
        - eps: Distance to detect eye movements. Adjust this parameter if results are not satisfying
        - kind: kind of interpolation in {'linear', 'cubic', 'quintic'}
    return:
        - move_indexes, blink_indexes, noise_indexes
    """
    x_pos    = np.array(eye_tracking[:,0])

    X        = np.stack((x_pos, np.linspace(0, len(x_pos), len(x_pos))*.5)).T
    clusters = cluster.dbscan(X, eps=eps, min_samples=5, metric='minkowski', p=2)
    move_indexes = np.where(clusters[1][1:] > clusters[1][:-1])[0] + 1

    noise_indexes = np.where(clusters[1] == -1)[0]
    blink_indexes = np.where(x_pos == 0)[0]

    return move_indexes, blink_indexes, noise_indexes

# Cell
def get_spherical_map(screen_pos, input_res=(281, 500), output_res=(360, 640), k_side=2, filling_pol="nonzero"):
    """
    Generate the mapping from normal image to an image projected on a spherical screen
    params:
        - screen_pos: Screen positions in shape (17, 10, (elev, azim)) obtained from calibration (from 80x80 pixel checker corners on a 1280x720px screen)
        - input_res: resolution of the input image
        - output_res: resolution of the output image
        - k_side: Kernel side's size to fill holes in the mapped image (2 -> 5*5 kernel)
        - filling_pol: Filling policy for the empty values, to which `f_fill` of `apply_spherical_map` will be applied
            nonzero: indexes of all nonzero elements
            closest: index of the closest nonzero element

    returns:
        - A mapping to be used in `apply_spherical_map`
    """
    assert filling_pol in ["nonzero", "closest"]

    screen_interp = interpolate_screen_pos(screen_pos, np.linspace(0,16, input_res[1], endpoint=True),
                                                       np.linspace(0, 9, input_res[0], endpoint=True))
    y_inres, x_inres = input_res
    y_res, x_res     = output_res
    xnew = np.linspace(screen_interp[:,:,1].min(), screen_interp[:,:,1].max(), x_res)
    ynew = np.linspace(screen_interp[:,:,0].min(), screen_interp[:,:,0].max(), y_res)
    map_img = np.zeros((y_res, x_res))

    y_map, x_map = np.empty(y_inres*x_inres, dtype=int), np.empty(y_inres*x_inres, dtype=int)
    for i, (y, x) in enumerate(zip(screen_interp[:,::-1,0].flatten(), screen_interp[:,::-1,1].flatten())):
        y_map[i] = np.argmin(ynew<y)
        x_map[i] = np.argmin(xnew<x)
    map_img[y_map, x_map] = 1

    y_nonzero, x_nonzero = np.nonzero(map_img==0) #Finds where the image is still zero
    fill_x_l, fill_y_l, nonzeros_l = [], [], []
    for y, x in zip(y_nonzero, x_nonzero):
        # Sets the limits to where to look for nonzeros pixels
        ylow, xlow = max(0, y-k_side), max(0, x-k_side)
        yhig, xhig = min(y+k_side+1, y_res), min(x+k_side+1, x_res)
        area = map_img[ylow:yhig, xlow:xhig]

        if np.any(area): #If there are pixels around
            fill_x_l.append(x)
            fill_y_l.append(y)
            nonz_y, nonz_x = np.nonzero(area)
            if filling_pol=="nonzero":
                nonzeros_l.append((nonz_y+ylow, nonz_x+xlow)) #store the nonzero slicing for later filling
            elif filling_pol=="closest":
                xx, yy = np.meshgrid(np.arange(xlow,xhig), np.arange(ylow,yhig))
                distances = np.sqrt((yy-y)**2+(xx-x)**2)
                idx_min   = np.argmin(distances[nonz_y, nonz_x])
                nonzeros_l.append(([nonz_y[idx_min]+ylow], [nonz_x[idx_min]+xlow]))

    return (y_map, x_map), (fill_y_l, fill_x_l, nonzeros_l)

def apply_spherical_map(img_src, mapping, output_res=(360, 640), f_fill=np.mean):
    """
    Apply a mapping to an input image
    params:
        - img_src: Source image to transform
        - mapping: Mapping obtained from `get_spherical_map`
        - output_res: Output resolution. Must match the parameter given to `get_spherical_map`
        - f_fill: Function to apply when filling the holes (e.g. np.median, np.mean)
    return:
        - The transformed image
    """
    (y_map, x_map), (fill_y_l, fill_x_l, nonzeros_l) = mapping

    transfo_img = np.zeros(output_res) + 128
    transfo_img[y_map, x_map] = img_src.flatten()

    for y, x, nonz in zip(fill_y_l, fill_x_l, nonzeros_l):
        transfo_img[y, x] = f_fill(transfo_img[nonz])
    return transfo_img