import numpy as np
from scipy.signal import correlate2d
from scipy.signal import fftconvolve
from scipy.ndimage.filters import sobel
import logging

def default_ray(n):

    y = np.zeros(n,dtype=np.int64)
    x = np.arange(n,dtype=np.int64)

    return np.vstack([y,x])

def rotate_ray(ray,theta):

    y,x = ray.astype(np.float64)

    xp = x*np.cos(theta) + y*np.sin(theta)
    yp = -x*np.sin(theta) + y*np.cos(theta)

    return np.vstack([yp.astype(np.int64),xp.astype(np.int64)])

def generate_rays(image_array, seed_pixel):

    N = 18 #200

    #mag, grad_x, grad_y = sobel_grad(image_array.astype('float'))

    shape = image_array.shape
    Y,X = np.mgrid[:shape[1],:shape[0]]

    n = int(np.sqrt(shape[0]**2 + shape[1]**2))
    angles = np.arange(N)*2.0*np.pi/N
    rays = []

    tangents = []

    ray_grads = []

    good_coords_mask = lambda y,x:  np.logical_and(np.logical_and(y>=0,y<shape[0]),np.logical_and(x>=0,x<shape[1]))

    for theta in angles:
        new_ray = rotate_ray(default_ray(n),theta)
        new_ray = new_ray.T + seed_pixel
        new_ray = new_ray.T


        mask = good_coords_mask(new_ray[0],new_ray[1])

        ym = new_ray[0][mask]
        xm = new_ray[1][mask]

        rays += [np.vstack([ym,xm])]

        t = np.array([np.sin(theta),np.cos(theta)])
        tangents += [t]

        #rg = t[1]*grad_x[ym,xm] + t[0]*grad_y[ym,xm]
        #rg[rg<0] = 0.0
        rg = image_array[ym,xm]
        #rg = rg[1:].astype(np.float64) - rg[:-1].astype(np.float64)
        # rg[rg<0] = 0.0
        ray_grads += [rg]


    return rays, ray_grads

def initial_pupil_point(image_array, bbox=None):
    """bbox is a tuple of (xmin, xmax, ymin, ymax)"""

    if bbox is not None:
        xmin, xmax, ymin, ymax = bbox
        crop_im = image_array[ymin:ymax, xmin:xmax]
    else:
        shape = image_array.shape
        crop_distance = 50
        crop_im = image_array[crop_distance:shape[0]-crop_distance, crop_distance:shape[1]-crop_distance]

    m = np.max(crop_im)

    dark_square = m*np.ones([30,30])
    #c = correlate2d(m-crop_im,dark_square,mode='same')
    c = fftconvolve(m-crop_im,dark_square[::-1,::-1],mode='same')
    y,x = np.where(c==np.max(c))

    if bbox is not None:
        ybar=int(np.mean(y))+ymin
        xbar=int(np.mean(x))+xmin
    else:
        ybar=int(np.mean(y))+crop_distance
        xbar=int(np.mean(x))+crop_distance

    return ybar, xbar

def initial_cr_point(image_array, bbox=None):
    """bbox is a tuple of (xmin, xmax, ymin, ymax)"""

    if bbox is not None:
        xmin, xmax, ymin, ymax = bbox
        crop_im = image_array[ymin:ymax, xmin:xmax]
    else:
        shape = image_array.shape
        crop_distance = 50
        crop_im = image_array[crop_distance:shape[0]-crop_distance, crop_distance:shape[1]-crop_distance]

    m = np.max(crop_im)
    mean = np.mean(crop_im)

    Y,X = np.meshgrid(np.arange(-20,20),np.arange(-20,20))
    bright_circle = np.zeros([40,40])
    mask = X**2 + Y**2 < 100.
    bright_circle[mask] = m
    bright_circle -= np.mean(bright_circle)

    #c = correlate2d(crop_im-mean,bright_circle,mode='same')
    c = fftconvolve(crop_im-mean,bright_circle[::-1,::-1],mode='same')
    y,x = np.where(c==np.max(c))

    if bbox is not None:
        ybar=int(np.mean(y))+ymin
        xbar=int(np.mean(x))+xmin
    else:
        ybar=int(np.mean(y))+crop_distance
        xbar=int(np.mean(x))+crop_distance

    return ybar, xbar

def sobel_grad(image_array):

    grad_y = sobel(image_array.astype(np.float64),0)
    grad_x = sobel(image_array.astype(np.float64),1)

    #print "grad_x dtype = ", grad_x.dtype

    mag = np.sqrt(grad_y**2 + grad_x**2) + 1e-16

    return mag, grad_x, grad_y

def medfilt_custom(x, kernel_size=3):
    '''This median filter returns 'nan' whenever any value in the kernal width is 'nan' and the median otherwise'''
    T = x.shape[0]
    delta = kernel_size/2

    x_med = np.zeros(x.shape)
    window = x[0:delta+1]
    if np.any(np.isnan(window)):
        x_med[0] = np.nan
    else:
        x_med[0] = np.median(window)

    # print window
    for t in range(1,T):
        window = x[t-delta:t+delta+1]
        # print window
        if np.any(np.isnan(window)):
            x_med[t] = np.nan
        else:
            x_med[t] = np.median(window)

    return x_med

def eccentricity(a1, a2):

    return np.sqrt(1.0 - (np.minimum(a1,a2)**2)/(np.maximum(a1,a2)**2))

def median_absolute_deviation(a, consistency_constant=1.4826):
    '''Calculate the median absolute deviation of a univariate dataset.

    Parameters
    ----------
    a : numpy.ndarray
        Sample data.
    consistency_constant : float
        Constant to make the MAD a consistent estimator of the population
        standard deviation (1.4826 for a normal distribution).

    Returns
    -------
    float
        Median absolute deviation of the data.
    '''
    return consistency_constant * np.nanmedian(np.abs(a - np.nanmedian(a)))

def post_process_cr(cr_params):
    """This will replace questionable values of the CR x and y position with 'nan'

        1)  threshold ellipse area by 99th percentile area distribution
        2)  median filter using custom median filter
        3)  remove deviations from discontinuous jumps

        The 'nan' values likely represent obscured CRs, secondary reflections, merges
        with the secondary reflection, or visual distortions due to the whisker or
        deformations of the eye"""

    area = np.pi*cr_params.T[3]*cr_params.T[4]

    # compute a threshold on the area of the cr ellipse
    dev = median_absolute_deviation(area)
    if dev == 0:
        logging.warning("Median absolute deviation is 0,"
                        "falling back to standard deviation.")
        dev = np.nanstd(area)
    threshold = np.nanmedian(area) + 3*dev

    x_center = cr_params.T[0]
    y_center = cr_params.T[1]

    # set x,y where area is over threshold to nan
    x_center[area>threshold] = np.nan
    y_center[area>threshold] = np.nan

    # median filter
    x_center_med = medfilt_custom(x_center, kernel_size=3)
    y_center_med = medfilt_custom(y_center, kernel_size=3)

    x_mask_finite = np.where(np.isfinite(x_center_med))[0]
    y_mask_finite = np.where(np.isfinite(y_center_med))[0]

    # if y increases discontinuously or x decreases discontinuously,
    #  that is probably a CR secondary reflection
    mean_x = np.mean(x_center_med[x_mask_finite])
    mean_y = np.mean(y_center_med[y_mask_finite])

    std_x = np.std(x_center_med[x_mask_finite])
    std_y = np.std(y_center_med[y_mask_finite])

    # set these extreme values to nan
    #x_center_med[x_center_med < mean_x - 3*std_x] = np.nan
    #y_center_med[y_center_med > mean_y + 3*std_y] = np.nan
    x_center_med[np.abs(x_center_med - mean_x) > 3*std_x] = np.nan
    y_center_med[np.abs(y_center_med - mean_y) > 3*std_y] = np.nan

    either_nan_mask = np.logical_and(np.isnan(x_center_med),np.isnan(y_center_med))
    x_center_med[either_nan_mask] = np.nan
    y_center_med[either_nan_mask] = np.nan

    new_cr = np.vstack([x_center_med, y_center_med]).T

    bad_points_mask = either_nan_mask

    return new_cr, bad_points_mask


def post_process_pupil(pupil_params):
    '''Filter pupil parameters to replace outliers with nan

    Parameters
    ----------
    pupil_params : numpy.ndarray
        (Nx5) array of pupil parameters [x, y, angle, axis1, axis2].

    Returns
    -------
    numpy.ndarray
        Pupil parameters with outliers replaced with nan
    '''
    area = np.pi*pupil_params.T[3]*pupil_params.T[4]
    threshold = np.percentile(area[np.isfinite(area)], 99)
    outlier_index = area > threshold
    pupil_params[outlier_index, :] = np.nan
    return pupil_params


def filter_bad_params(params, frame_width, frame_height):
    '''Replace positions outside image with nan'''
    params[(params[:,0] > frame_width) | (params[:,0] < 0), :] = np.nan
    params[(params[:,1] > frame_height) | (params[:,1] < 0), :] = np.nan
    return params
