import numpy as np

def spatial_rate_map(x, y, t, spike_train, binsize=0.01, box_xlen=1,
                     box_ylen=1, mask_unvisited=True, convolve=True,
                     return_bins=False, smoothing=0.02):
    """Divide a 2D space in bins of size binsize**2, count the number of spikes
    in each bin and divide by the time spent in respective bins. The map can
    then be convolved with a gaussian kernel of size csize determined by the
    smoothing factor, binsize and box_xlen.

    Parameters
    ----------
    spike_train : neo.SpikeTrain
    x : float
        1d vector of x positions
    y : float
        1d vector of y positions
    t : float
        1d vector of times at x, y positions
    binsize : float
        spatial binsize
    box_xlen : quantities scalar in m
        side length of quadratic box
    mask_unvisited: bool
        mask bins which has not been visited by nans
    convolve : bool
        convolve the rate map with a 2D Gaussian kernel

    Returns
    -------
    out : rate map
    if return_bins = True
    out : rate map, xbins, ybins
    """
    if not all([len(var) == len(var2) for var in [x,y,t] for var2 in [x,y,t]]):
        raise ValueError('x, y, t must have same number of elements')
    if box_xlen < x.max() or box_ylen < y.max():
        raise ValueError('box length must be larger or equal to max path length')
    from decimal import Decimal as dec
    decimals = 1e10
    remainderx = dec(float(box_xlen)*decimals) % dec(float(binsize)*decimals)
    remaindery = dec(float(box_ylen)*decimals) % dec(float(binsize)*decimals)
    if remainderx != 0 or remaindery != 0:
        raise ValueError('the remainder should be zero i.e. the ' +
                         'box length should be an exact multiple ' +
                         'of the binsize')

    # interpolate one extra timepoint
    t_ = np.append(t, t[-1] + np.median(np.diff(t)))
    spikes_in_bin, _ = np.histogram(spike_train, t_)
    time_in_bin = np.diff(t_)
    xbins = np.arange(0, box_xlen + binsize, binsize)
    ybins = np.arange(0, box_ylen + binsize, binsize)
    ix = np.digitize(x, xbins, right=True)
    iy = np.digitize(y, ybins, right=True)
    spike_pos = np.zeros((xbins.size, ybins.size))
    time_pos = np.zeros((xbins.size, ybins.size))
    for n in range(len(x)):
        spike_pos[ix[n], iy[n]] += spikes_in_bin[n]
        time_pos[ix[n], iy[n]] += time_in_bin[n]
    # correct for shifting of map
    spike_pos = spike_pos[1:, 1:]
    time_pos = time_pos[1:, 1:]
    with np.errstate(divide='ignore', invalid='ignore'):
        rate = np.divide(spike_pos, time_pos)
    if convolve:
        rate[np.isnan(rate)] = 0.  # for convolution
        from astropy.convolution import Gaussian2DKernel, convolve_fft
        csize = (box_xlen / binsize) * smoothing
        kernel = Gaussian2DKernel(csize)
        rate = convolve_fft(rate, kernel)  # TODO edge correction
    if mask_unvisited:
        was_in_bin = np.asarray(time_pos, dtype=bool)
        rate[np.invert(was_in_bin)] = np.nan
    if return_bins:
        return rate.T, xbins, ybins
    else:
        return rate.T

def gridness(rate_map, box_xlen, box_ylen, return_acorr=False,
             step_size=0.1, method='iter', return_masked_acorr=False):
    '''Calculates gridness of a rate map. Calculates the normalized
    autocorrelation (A) of a rate map B where A is given as
    A = 1/n\Sum_{x,y}(B - \bar{B})^{2}/\sigma_{B}^{2}. Further, the Pearsson's
    product-moment correlation coefficients is calculated between A and A_{rot}
    rotated 30 and 60 degrees. Finally the gridness is calculated as the
    difference between the minimum of coefficients at 60 degrees and the
    maximum of coefficients at 30 degrees i.e. gridness = min(r60) - max(r30).

    If the method 'iter' is chosen:
    In order to focus the analysis on symmetry of A the the central and the
    outer part of the gridness is maximized by increasingly mask A at steps of
    ``step_size``.

    If the method 'puncture' is chosen:
    This is the standard way of calculating gridness, by masking the central
    autocorrelation bump, in addition to rounding the map. See examples.

    Parameters
    ----------
    rate_map : numpy.ndarray
    box_xlen : float
        side length of quadratic box
    step_size : float
        step size in masking, only applies to the method "iter"
    return_acorr : bool
        return autocorrelation map or not
    return_masked_acorr : bool
        return masked autocorrelation map or not
    method : 'iter' or 'puncture'

    Returns
    -------
    out : gridness, (autocorrelation map, masked autocorrelation map)

    Examples
    --------
    >>> from exana.tracking.tools import make_test_grid_rate_map
    >>> import matplotlib.pyplot as plt
    >>> rate_map, pos = make_test_grid_rate_map()
    >>> iter_score = gridness(rate_map, box_xlen=1, box_ylen=1, method='iter')
    >>> print('%.2f' % iter_score)
    1.39
    >>> puncture_score = gridness(rate_map, box_xlen=1, box_ylen=1, method='puncture')
    >>> print('%.2f' % puncture_score)
    0.96

    .. plot::

        import matplotlib.pyplot as plt
        import numpy as np
        from exana.tracking.tools import make_test_grid_rate_map
        from exana.tracking import gridness
        import matplotlib.pyplot as plt
        rate_map, _ = make_test_grid_rate_map()
        fig, axs = plt.subplots(2, 2)
        g1, acorr, m_acorr1 = gridness(rate_map, box_xlen=1,
                                         box_ylen=1, return_acorr=True,
                                         return_masked_acorr=True,
                                         method='iter')
        g2, m_acorr2 = gridness(rate_map, box_xlen=1,
                                         box_ylen=1,
                                         return_masked_acorr=True,
                                         method='puncture')
        mats = [rate_map, m_acorr1, acorr, m_acorr2]
        titles = ['Rate map', 'Masked acorr "iter", gridness = %.2f' % g1,
                  'Autocorrelation',
                  'Masked acorr "puncture", gridness = %.2f' % g2]
        for ax, mat, title in zip(axs.ravel(), mats, titles):
            ax.imshow(mat)
            ax.set_title(title)
        plt.tight_layout()
        plt.show()
    '''
    import numpy.ma as ma
    from exana.misc.tools import fftcorrelate2d
    from exana.tracking.tools import gaussian2D
    from scipy.optimize import curve_fit

    tmp_map = rate_map.copy()
    tmp_map[~np.isfinite(tmp_map)] = 0
    acorr = fftcorrelate2d(tmp_map, tmp_map, mode='full', normalize=True)
    rows, cols = acorr.shape
    b_x = np.linspace(- box_xlen / 2., box_xlen / 2., rows)
    b_y = np.linspace(- box_ylen / 2., box_ylen / 2., cols)
    B_x, B_y = np.meshgrid(b_x, b_y)
    if method == 'iter':
        if return_masked_acorr: m_acorrs = []
        gridscores = []
        for outer in np.arange(box_xlen / 4, box_xlen / 2, step_size):
            m_acorr = ma.masked_array(
                acorr, mask=np.sqrt(B_x**2 + B_y**2) > outer)
            for inner in np.arange(0, box_xlen / 4, step_size):
                m_acorr = ma.masked_array(
                    m_acorr, mask=np.sqrt(B_x**2 + B_y**2) < inner)
                r30, r60 = rotate_corr(m_acorr)
                gridscores.append(np.min(r60) - np.max(r30))
                if return_masked_acorr: m_acorrs.append(m_acorr)
        gridscore = max(gridscores)
        if return_masked_acorr: m_acorr = m_acorrs[gridscores.index(gridscore)]
    elif method == 'puncture':
        # round picture edges
        _gaussian = lambda pos, a, s: gaussian2D(a, pos[0], pos[1], 0, 0, s).ravel()
        p0 = (max(acorr.ravel()), min(box_xlen, box_ylen) / 100)
        popt, pcov = curve_fit(_gaussian, (B_x, B_y), acorr.ravel(), p0=p0)
        m_acorr = ma.masked_array(
            acorr, mask=np.sqrt(B_x**2 + B_y**2) > min(box_xlen, box_ylen) / 2)
        m_acorr = ma.masked_array(
            m_acorr, mask=np.sqrt(B_x**2 + B_y**2) < popt[1])
        r30, r60 = rotate_corr(m_acorr)
        gridscore = float(np.min(r60) - np.max(r30))
    if return_acorr and return_masked_acorr:
        return gridscore, acorr, m_acorr
    if return_masked_acorr:
        return gridscore, m_acorr
    if return_acorr:
        return gridscore, acorr  # acorrs[grids.index(max(grids))]
    else:
        return gridscore

def rotate_corr(acorr):
    from exana.misc.tools import masked_corrcoef2d
    from scipy.ndimage.interpolation import rotate
    angles = range(30, 180+30, 30)
    corr = []
    # Rotate and compute correlation coefficient
    for angle in angles:
        rot_acorr = rotate(acorr, angle, reshape=False)
        corr.append(masked_corrcoef2d(rot_acorr, acorr)[0, 1])
    r60 = corr[1::2]
    r30 = corr[::2]
    return r30, r60

def occupancy_map(x, y, t,
                  binsize=0.01,
                  box_xlen=1,
                  box_ylen=1,
                  mask_unvisited=True,
                  convolve=True,
                  return_bins=False,
                  smoothing=0.02):
    '''Divide a 2D space in bins of size binsize**2, count the time spent
    in each bin. The map can  be convolved with a gaussian kernel of size
    csize determined by the smoothing factor, binsize and box_xlen.

    Parameters
    ----------
    x : array
        1d vector of x positions
    y : array
        1d vector of y positions
    t : array
        1d vector of times at x, y positions
    binsize : float
        spatial binsize
    box_xlen : float
        side length of quadratic box
    mask_unvisited: bool
        mask bins which has not been visited by nans
    convolve : bool
        convolve the rate map with a 2D Gaussian kernel

    Returns
    -------
    occupancy_map : numpy.ndarray
    if return_bins = True
    out : occupancy_map, xbins, ybins
    '''

    if not all([len(var) == len(var2) for var in [
            x, y, t] for var2 in [x, y, t]]):
        raise ValueError('x, y, t must have same number of elements')
    if box_xlen < x.max() or box_ylen < y.max():
        raise ValueError(
            'box length must be larger or equal to max path length')
    from decimal import Decimal as dec
    decimals = 1e10
    remainderx = dec(float(box_xlen)*decimals) % dec(float(binsize)*decimals)
    remaindery = dec(float(box_ylen)*decimals) % dec(float(binsize)*decimals)
    if remainderx != 0 or remaindery != 0:
        raise ValueError('the remainder should be zero i.e. the ' +
                         'box length should be an exact multiple ' +
                         'of the binsize')

    # interpolate one extra timepoint
    t_ = np.array(t.tolist() + [t.max() + np.median(np.diff(t))])
    time_in_bin = np.diff(t_)
    xbins = np.arange(0, box_xlen + binsize, binsize)
    ybins = np.arange(0, box_ylen + binsize, binsize)
    ix = np.digitize(x, xbins, right=True)
    iy = np.digitize(y, ybins, right=True)
    time_pos = np.zeros((xbins.size, ybins.size))
    for n in range(len(x) - 1):
        time_pos[ix[n], iy[n]] += time_in_bin[n]
    # correct for shifting of map since digitize returns values at right edges
    time_pos = time_pos[1:, 1:]
    if convolve:
        rate[np.isnan(rate)] = 0.  # for convolution
        from astropy.convolution import Gaussian2DKernel, convolve_fft
        csize = (box_xlen / binsize) * smoothing
        kernel = Gaussian2DKernel(csize)
        rate = convolve_fft(rate, kernel)  # TODO edge correction
    if mask_unvisited:
        was_in_bin = np.asarray(time_pos, dtype=bool)
        rate[np.invert(was_in_bin)] = np.nan
    if return_bins:
        return rate.T, xbins, ybins
    else:
        return rate.T

def nvisits_map(x, y, t,
                binsize=0.01,
                box_xlen=1,
                box_ylen=1,
                return_bins=False):
    '''Divide a 2D space in bins of size binsize**2, count the
    number of visits in each bin. The map can  be convolved with
    a gaussian kernel of size  determined by the smoothing factor,
    binsize and box_xlen.

    Parameters
    ----------
    x : array
        1d vector of x positions
    y : array
        1d vector of y positions
    t : array
        1d vector of times at x, y positions
    binsize : float
        spatial binsize
    box_xlen : float
        side length of quadratic box

    Returns
    -------
    nvisits_map : numpy.ndarray
    if return_bins = True
    out : nvisits_map, xbins, ybins
    '''

    if not all([len(var) == len(var2) for var in [
            x, y, t] for var2 in [x, y, t]]):
        raise ValueError('x, y, t must have same number of elements')
    if box_xlen < x.max() or box_ylen < y.max():
        raise ValueError(
            'box length must be larger or equal to max path length')
    from decimal import Decimal as dec
    decimals = 1e10
    remainderx = dec(float(box_xlen)*decimals) % dec(float(binsize)*decimals)
    remaindery = dec(float(box_ylen)*decimals) % dec(float(binsize)*decimals)
    if remainderx != 0 or remaindery != 0:
        raise ValueError('the remainder should be zero i.e. the ' +
                         'box length should be an exact multiple ' +
                         'of the binsize')

    xbins = np.arange(0, box_xlen + binsize, binsize)
    ybins = np.arange(0, box_ylen + binsize, binsize)
    ix = np.digitize(x, xbins, right=True)
    iy = np.digitize(y, ybins, right=True)
    nvisits_map = np.zeros((xbins.size, ybins.size))
    for n in range(len(x)):
        if n == 0:
            nvisits_map[ix[n], iy[n]] = 1
        else:
            if ix[n-1] != ix[n] or iy[n-1] != iy[n]:
                nvisits_map[ix[n], iy[n]] += 1
    # correct for shifting of map since digitize returns values at right edges
    nvisits_map = nvisits_map[1:, 1:]
    if return_bins:
        return nvisits_map.T, xbins, ybins
    else:
        return nvisits_map.T

def spatial_rate_map_1d(x, t, spike_train,
                        binsize=0.01,
                        track_len=1,
                        mask_unvisited=True,
                        convolve=True,
                        return_bins=False,
                        smoothing=0.02):
    """Take x coordinates of linear track data, divide in bins of binsize,
    count the number of spikes  in each bin and  divide by the time spent in
    respective bins. The map can then be convolved with a gaussian kernel of
    size csize determined by the smoothing factor, binsize and box_xlen.

    Parameters
    ----------
    spike_train : array
    x : array
        1d vector of x positions
    t : array
        1d vector of times at x, y positions
    binsize : float
        spatial binsize
    box_xlen : float
        side length of quadratic box
    mask_unvisited: bool
        mask bins which has not been visited by nans
    convolve : bool
        convolve the rate map with a 2D Gaussian kernel

    Returns
    -------
    out : rate map
    if return_bins = True
    out : rate map, xbins
    """
    if not all([len(var) == len(var2) for var in [x, t] for var2 in [x, t]]):
        raise ValueError('x, t must have same number of elements')
    if track_len < x.max():
        raise ValueError('track length must be\
        larger or equal to max path length')
    from decimal import Decimal as dec
    decimals = 1e10
    remainderx = dec(float(track_len)*decimals) % dec(float(binsize)*decimals)
    if remainderx != 0:
        raise ValueError('the remainder should be zero i.e. the ' +
                         'box length should be an exact multiple ' +
                         'of the binsize')
    # interpolate one extra timepoint
    t_ = np.array(t.tolist() + [t.max() + np.median(np.diff(t))])
    spikes_in_bin, _ = np.histogram(spike_train, t_)
    time_in_bin = np.diff(t_)
    xbins = np.arange(0, track_len + binsize, binsize)
    ix = np.digitize(x, xbins, right=True)
    spike_pos = np.zeros(xbins.size)
    time_pos = np.zeros(xbins.size)
    for n in range(len(x)):
        spike_pos[ix[n]] += spikes_in_bin[n]
        time_pos[ix[n]] += time_in_bin[n]
    # correct for shifting of map since digitize returns values at right edges
    spike_pos = spike_pos[1:]
    time_pos = time_pos[1:]
    with np.errstate(divide='ignore', invalid='ignore'):
        rate = np.divide(spike_pos, time_pos)
    if convolve:
        rate[np.isnan(rate)] = 0.  # for convolution
        from astropy.convolution import Gaussian2DKernel, convolve_fft
        csize = (track_len / binsize) * smoothing
        kernel = Gaussian2DKernel(csize)
        rate = convolve_fft(rate, kernel)  # TODO edge correction
    if mask_unvisited:
        was_in_bin = np.asarray(time_pos, dtype=bool)
        rate[np.invert(was_in_bin)] = np.nan
    if return_bins:
        return rate.T, xbins
    else:
        return rate.T

def separate_fields(rate_map, laplace_thrsh=0, center_method='maxima',
        cutoff_method='none', box_xlen=1, box_ylen=1, index=False):
    """Separates fields using the laplacian to identify fields separated by
    a negative second derivative.

    Parameters
    ----------
    rate_map : np 2d array
        firing rate in each bin
    laplace_thrsh : float
        value of laplacian to separate fields by relative to the minima. Should be
        on the interval 0 to 1, where 0 cuts off at 0 and 1 cuts off at
        min(laplace(rate_map)). Default 0.
    center_method : string
        method to find field centers. Valid options = ['center_of_mass',
        'maxima','gaussian_fit']
    cutoff_method (optional) : string or function
        function to exclude small fields. If local field value of function
        is lower than global function value, the field is excluded. Valid
        string_options = ['median', 'mean','none'].
    index : bool, default False
        return bump center values as index or xy-pos

    Returns
    -------
    fields : numpy array, shape like rate_map.
        contains areas all filled with same value, corresponding to fields
        in rate_map. The values are in range(1,nFields + 1), sorted by size of the
        field (sum of all field values). 0 elsewhere.
    n_field : int
        field count
    bump_centers : (n_field x 2) np ndarray
        Coordinates of field centers
    """

    cutoff_functions = {'mean':np.mean, 'median':np.median, 'none':None}
    if not callable(cutoff_method):
        try:
            cutoff_func = cutoff_functions[cutoff_method]
        except KeyError:
            msg = "invalid cutoff_method flag '%s'" % cutoff_method
            raise ValueError(msg)
    else:
        cutoff_func = cutoff_method

    from scipy import ndimage

    l = ndimage.laplace(rate_map)

    l[l>laplace_thrsh*np.min(l)] = 0

    # Labels areas of the laplacian not connected by values > 0.
    fields, n_fields = ndimage.label(l)

    # index 0 is the background
    indx = np.arange(1,n_fields+1)

    # Use cutoff method to remove unwanted fields
    if cutoff_method != 'none':
        try:
            total_value = cutoff_func(fields)
        except:
            print('Unexpected error, cutoff_func doesnt like the input:')
            raise

        field_values = ndimage.labeled_comprehension(rate_map, fields, indx,
                cutoff_func, float, 0)
        try:
            is_field = field_values >= total_value
        except:
            print('cutoff_func return_values doesnt want to compare:')
            raise

        if np.sum(is_field) == 0:
            return np.zeros(rate_map.shape), 0, np.array([[],[]])

        for i in indx:
            if not is_field[i-1]:
                fields[fields == i] = 0

        n_fields = ndimage.label(fields, output=fields)
        indx = np.arange(1,n_fields + 1)

    # Sort by largest mean
    sizes = ndimage.labeled_comprehension(rate_map, fields, indx,
            np.mean, float, 0)
    size_sort = np.argsort(sizes)[::-1]
    new = np.zeros_like(fields)
    for i in np.arange(n_fields):
        new[fields == size_sort[i]+1] = i+1
    fields = new

    bc = get_bump_centers(rate_map,labels=fields,ret_index=index,indices=indx,method=center_method,
                          units=box_xlen.units)

    # TODO exclude fields where maxima is on the edge of the field?
    return fields, n_fields, bc

def get_bump_centers(rate_map, labels, ret_index=False, indices=None, method='maxima',
        units=1):
    """Finds center of fields at labels."""

    from scipy import ndimage

    if method not in ['maxima','center_of_mass','gaussian_fit']:
        msg = "invalid center_method flag '%s'" % method
        raise ValueError(msg)
    if indices is None:
        indices = np.arange(1,np.max(labels)+1)
    if method == 'maxima':
        bc = ndimage.maximum_position(rate_map, labels=labels,
                index=indices)
    elif method == 'center_of_mass':
        bc = ndimage.center_of_mass(rate_map, labels=labels, index=indices)
    elif method == 'gaussian_fit':
        from  exana.tracking.tools import fit_gauss_asym
        bc = np.zeros((len(indices),2))
        import matplotlib.pyplot as plt
        for i in indices:
            r = rate_map.copy()
            r[labels != i] = 0
            popt = fit_gauss_asym(r, return_data=False)
            # TODO Find out which axis is x and which is y
            bc[i-1] = (popt[2],popt[1])
        if ret_index:
            msg = 'ret_index not implemented for gaussian fit'
            raise NotImplementedError(msg)
    if not ret_index and not method=='gaussian_fit':
        bc = (bc + np.array((0.5,0.5)))/rate_map.shape
    return np.array(bc)*units

def find_avg_dist(rate_map, thrsh = 0, plot=False):
    """Uses autocorrelation and separate_fields to find average distance
    between bumps. Is dependent on high gridness to get separate bumps in
    the autocorrelation

    Parameters
    ----------
    rate_map : np 2d array
               firing rate in each bin

    thrsh (optional) : float, default 0
        cutoff value for the laplacian of the autocorrelation function.
        Should be a negative number. Gives better separation if bumps are
        connected by "bridges" or saddles where the laplacian is negative.
    plot (optional) : bool, default False
        plot acorr and the separated acorr, with bump centers
    Returns
    -------
    avg_dist : float
        relative units from 0 to 1 of the box size
        """

    from scipy.ndimage import maximum_position
    from exana.misc.tools import fftcorrelate2d

    # autocorrelate. Returns array (2x - 1) the size of rate_map
    acorr = fftcorrelate2d(rate_map,rate_map, mode = 'full', normalize = True)

    #acorr[acorr<0] = 0 # TODO Fix this
    f, nf, bump_centers = separate_fields(acorr,laplace_thrsh=thrsh,
            center_method='maxima',cutoff_method='median')
                                         # TODO Find a way to find valid value for
                                         # thrsh, or remove.

    bump_centers = np.array(bump_centers)

    # find dists from center in (autocorrelation)relative units (from 0 to 1)
    distances = np.linalg.norm(bump_centers - (0.5,0.5), axis = 1)

    dist_sort = np.argsort(distances)
    distances = distances[dist_sort]

    # use maximum 6 closest values except center value
    avg_dist = np.median(distances[1:7])

    # correct for difference in shapes
    avg_dist *= acorr.shape[0]/rate_map.shape[0] # = 1.98

    # TODO : raise warning if too big difference between points
    if plot:
        import matplotlib.pyplot as plt
        fig,[ax1,ax2] = plt.subplots(1,2)

        ax1.imshow(acorr,extent  = (0,1,0,1),origin='lower')
        ax1.scatter(*(bump_centers[:,::-1].T))
        ax2.imshow(f,extent  = (0,1,0,1),origin='lower')
        ax2.scatter(*(bump_centers[:,::-1].T))
    return avg_dist

def fit_hex(bump_centers, avg_dist=None, plot_bumps = False, method='best'):
    """Fits a hex grid to a given set of bumps. Uses the three bumps most

    Parameters
    ----------
    bump_centers : Nx2 np.array
                x,y positions of bump centers, x,y /in (0,1)

    avg_dist (optional): float
                average spacing between bumps

    plot_bumps (optional): bool
                if True, plots at the three bumps most likely to be in
                correct hex-position to the current matplotlib axes.

    method (optional): string, valid options: ['closest', 'best']
                method to find angle from neighboring bumps.
                'closest' uses six bumps nearest to center bump
                'best' uses the two bumps nearest to avg_dist

    Returns
    -------
    displacement : float
                   distance of bump closest to the center in meters
    orientation : float
                  orientation of hexagon (in degrees)
    """

    valid_methods = ['closest', 'best']
    if method not in valid_methods:
        msg = "invalid method flag '%s'" % method
        raise ValueError(msg)
    bump_centers = np.array(bump_centers)

    # sort by distance to center
    d = np.linalg.norm(bump_centers - (0.5,0.5), axis=1)
    d_sort = np.argsort(d)
    dist_sorted = bump_centers[d_sort]
    center_bump = dist_sorted[0]; others = dist_sorted[1:]

    displacement = d[d_sort][0]

    # others distances to center bumps
    relpos = others - center_bump
    reldist = np.linalg.norm(relpos, axis=1)

    if method == 'closest':
        # get 6 closest bumps
        rel_sort = np.argsort(reldist)
        closest = others[rel_sort][:6]
        relpos = relpos[rel_sort][:6]
    elif method == 'best':
        # get 2 bumps such that /sum_{i\neqj}(\abs{r_i-r_j}-avg_ist)^2 is minimized
        squares = 1e32*np.ones((others.shape[0], others.shape[0]))

        for i in range(len(relpos)):
            for j in range(i,len(relpos)):
                rel1 = (reldist[i] - avg_dist)**2
                rel2 = (reldist[j] - avg_dist)**2
                rel3 = (np.linalg.norm(relpos[i]-relpos[j]) - avg_dist)**2
                squares[i,j] = rel1 + rel2 + rel3
        rel_slice = np.unravel_index(np.argmin(squares), squares.shape)
        rel_slice = np.array(rel_slice)
        #rel_sort = np.argsort(np.abs(reldist-avg_dist))
        closest = others[rel_slice]
        relpos = relpos[rel_slice]

    # sort by angle
    a = np.arctan2(relpos[:,1], relpos[:,0])%(2*np.pi)
    a_sort = np.argsort(a)

    # extract lowest angle and convert to degrees
    orientation = a[a_sort][0] *180/np.pi

    # hex grid is symmetric under rotations of 60deg
    orientation %= 60

    if plot_bumps:
        import matplotlib.pyplot as plt
        ax=plt.gca()
        i = 1
        xmin, xmax = ax.get_xlim()
        ymin, ymax = ax.get_ylim()
        dx = xmax-xmin; dy = ymax - ymin

        closest = closest[a_sort]

        edges = [center_bump] if method == 'best' else []
        edges += [c for c in closest]
        edges = np.array(edges)*(dx,dy) + (xmin, ymin)
        poly = plt.Polygon(edges, alpha=0.5,color='r')
        ax.add_artist(poly)
    return displacement, orientation

def calculate_grid_geometry(rate_map, plot_fields=False, **kwargs):
    """Calculates quantitative information about grid field.
    Find bump centers, bump spacing, center diplacement and hexagon
    orientation

    Parameters
    ----------
    rate_map : np 2d array
        firing rate in each bin
    plot_fields : if True, plots the field labels with field centers to the
        current matplotlib ax. Default False
    thrsh : float, default 0
        see find_avg_dist()
    center_method : string, valid options: ['maxima', 'center_of_mass']
        default: 'center_of_mass'
        see separate_fields()
    method : string, valid options: ['closest', 'best']
        see fit_hex()

    Returns
    -------
    bump_centers : 2d np.array
        x,y positions of bump centers
    avg_dist : float
        average spacing between bumps, \in [0,1]
    displacement : float
        distance of bump closest to the center
    orientation : float
        orientation of hexagon (in degrees)

    Examples
    --------
    >>> import numpy as np
    >>> rate_map = np.zeros((5,5))
    >>> pos = np.array([  [0,2],
    ...                [1,0],[1,4],
    ...                   [2,2],
    ...                [3,0],[3,4],
    ...                   [4,2]])
    >>> for(i,j) in pos:
    ...     rate_map[i,j] = 1
    ...
    >>> result = calculate_grid_geometry(rate_map)
    """

    # TODO add back the following when it is correct
    # (array([[0.5, 0.9],
           # [0.9, 0.7],
           # [0.1, 0.7],
           # [0.5, 0.5],
           # [0.9, 0.3],
           # [0.1, 0.3],
           # [0.5, 0.1]]) * m, 0.4472135954999579, 0.0, 26.565051177077983)

    from scipy.ndimage import mean, center_of_mass

    # TODO: smooth data?
    # smooth_rate_map = lambda x:x
    # rate_map = smooth_rate_map(rate_map)

    center_method = kwargs.pop('center_method',None)
    if center_method:
        fields, nfields, bump_centers = separate_fields(rate_map,
                                        center_method=center_method)
    else:
        fields, nfields, bump_centers = separate_fields(rate_map)

    if bump_centers.size == 0:
        import warnings
        msg = 'couldnt find bump centers, returning None'
        warnings.warn(msg, RuntimeWarning, stacklevel=2)
        return None,None,None,None,

    sh = np.array(rate_map.shape)

    if plot_fields:
        print(fields)
        import matplotlib.pyplot as plt
        x=np.linspace(0,1,sh[0]+1)
        y=np.linspace(0,1,sh[1]+1)
        x,y = np.meshgrid(x,y)
        ax = plt.gca()
        print('nfields: ',nfields)
        plt.pcolormesh(x,y, fields)

    # switch from row-column to x-y
    bump_centers = bump_centers[:,::-1]

    thrsh = kwargs.pop('thrsh', None)
    if thrsh:
        avg_dist = find_avg_dist(rate_map, thrsh)
    else:
        avg_dist = find_avg_dist(rate_map)

    displacement, orientation = fit_hex(bump_centers, avg_dist,
            plot_bumps=plot_fields, **kwargs)

    return bump_centers, avg_dist, displacement, orientation

class RandomDisplacementBounds(object):
    """random displacement with bounds"""
    def __init__(self, xmin, xmax, stepsize=0.5):
        self.xmin = np.array(xmin)
        self.xmax = np.array(xmax)
        self.stepsize = stepsize

    def __call__(self, x):
        """take a random step but ensure the new position is within the bounds"""
        while True:

            # this could be done in a much more clever way, but it will work for example purposes
            xnew = x + (self.xmax-self.xmin)*np.random.uniform(-self.stepsize,
                                                               self.stepsize, np.shape(x))
            if np.all(xnew < self.xmax) and np.all(xnew > self.xmin):
                break
        return xnew

def optimize_sep_fields(rate_map,step = 0.04, niter=40, T = 1.0, method = 'SLSQP',
        glob=True, x0 = [0.065,0.1],callback=None):
    """Optimizes the separation of the fields by minimizing an error
    function

    Parameters
    ----------
    rate_map :
    method :
        valid methods=['L-BFGS-B', 'TNC', 'SLSQP']
    x0 : list
        initial values for smoothing smoothing and laplace_thrsh

    Returns
    --------
    res :
        Result of the optimization. Contains smoothing and laplace_thrsh in
        attribute res.x"""

    from scipy import optimize
    from exana.tracking.tools import separation_error_func as err_func

    valid_methods = ['L-BFGS-B', 'TNC', 'SLSQP']
    if method not in valid_methods:
        raise ValueError('invalid method flag %s' %method)

    rate_map[np.isnan(rate_map)] = 0.

    method = 'SLSQP'
    xmin = [0.025, 0]
    xmax = [0.2,  1]
    bounds = [(low,high) for low,high in zip(xmin,xmax)]

    obj_func = lambda args: err_func(args[0], args[1], rate_map)

    if glob:
        take_step = RandomDisplacementBounds(xmin, xmax,stepsize=step)
        minimizer_kwargs = dict(method=method, bounds=bounds)
        res = optimize.basinhopping(obj_func, x0, niter=niter, T = T,
                minimizer_kwargs=minimizer_kwargs,
                take_step=take_step,callback=callback)
    else:
        res = optimize.minimize(obj_func, x0, method=method, bounds = bounds, options={'disp': True})
    return res

if __name__ == "__main__":
    import doctest
    doctest.testmod()


"""
The :mod:`sklearn.pipeline` module implements utilities to build a composite
estimator, as a chain of transforms and estimators.
"""
# Author: Edouard Duchesnay
#         Gael Varoquaux
#         Virgile Fritsch
#         Alexandre Gramfort
#         Lars Buitinck
# Licence: BSD

from collections import defaultdict

import numpy as np
from scipy import sparse

from .base import BaseEstimator, TransformerMixin
from .externals.joblib import Parallel, delayed
from .externals import six
from .utils import tosequence
from .utils.metaestimators import if_delegate_has_method
from .externals.six import iteritems

__all__ = ['Pipeline', 'FeatureUnion']

class Pipeline(BaseEstimator):
    """Pipeline of transforms with a final estimator.

    Sequentially apply a list of transforms and a final estimator.
    Intermediate steps of the pipeline must be 'transforms', that is, they
    must implement fit and transform methods.
    The final estimator only needs to implement fit.

    The purpose of the pipeline is to assemble several steps that can be
    cross-validated together while setting different parameters.
    For this, it enables setting parameters of the various steps using their
    names and the parameter name separated by a '__', as in the example below.

    Read more in the :ref:`User Guide <pipeline>`.

    Parameters
    ----------
    steps : list
        List of (name, transform) tuples (implementing fit/transform) that are
        chained, in the order in which they are chained, with the last object
        an estimator.

    Attributes
    ----------
    named_steps : dict
        Read-only attribute to access any step parameter by user given name.
        Keys are step names and values are steps parameters.

    Examples
    --------
    >>> from sklearn import svm
    >>> from sklearn.datasets import samples_generator
    >>> from sklearn.feature_selection import SelectKBest
    >>> from sklearn.feature_selection import f_regression
    >>> from sklearn.pipeline import Pipeline
    >>> # generate some data to play with
    >>> X, y = samples_generator.make_classification(
    ...     n_informative=5, n_redundant=0, random_state=42)
    >>> # ANOVA SVM-C
    >>> anova_filter = SelectKBest(f_regression, k=5)
    >>> clf = svm.SVC(kernel='linear')
    >>> anova_svm = Pipeline([('anova', anova_filter), ('svc', clf)])
    >>> # You can set the parameters using the names issued
    >>> # For instance, fit using a k of 10 in the SelectKBest
    >>> # and a parameter 'C' of the svm
    >>> anova_svm.set_params(anova__k=10, svc__C=.1).fit(X, y)
    ...                                              # doctest: +ELLIPSIS
    Pipeline(steps=[...])
    >>> prediction = anova_svm.predict(X)
    >>> anova_svm.score(X, y)                        # doctest: +ELLIPSIS
    0.77...
    >>> # getting the selected features chosen by anova_filter
    >>> anova_svm.named_steps['anova'].get_support()
    ... # doctest: +NORMALIZE_WHITESPACE
    array([ True,  True,  True, False, False,  True, False,  True,  True, True,
           False, False,  True, False,  True, False, False, False, False,
           True], dtype=bool)
    """

    # BaseEstimator interface

    def __init__(self, steps):
        names, estimators = zip(*steps)
        if len(dict(steps)) != len(steps):
            raise ValueError("Provided step names are not unique: %s" % (names,))

        # shallow copy of steps
        self.steps = tosequence(steps)
        transforms = estimators[:-1]
        estimator = estimators[-1]

        for t in transforms:
            if (not (hasattr(t, "fit") or hasattr(t, "fit_transform")) or not
                    hasattr(t, "transform")):
                raise TypeError("All intermediate steps of the chain should "
                                "be transforms and implement fit and transform"
                                " '%s' (type %s) doesn't)" % (t, type(t)))

        if not hasattr(estimator, "fit"):
            raise TypeError("Last step of chain should implement fit "
                            "'%s' (type %s) doesn't)"
                            % (estimator, type(estimator)))

    @property
    def _estimator_type(self):
        return self.steps[-1][1]._estimator_type

    def get_params(self, deep=True):
        if not deep:
            return super(Pipeline, self).get_params(deep=False)
        else:
            out = self.named_steps
            for name, step in six.iteritems(self.named_steps):
                for key, value in six.iteritems(step.get_params(deep=True)):
                    out['%s__%s' % (name, key)] = value

            out.update(super(Pipeline, self).get_params(deep=False))
            return out

    @property
    def named_steps(self):
        return dict(self.steps)

    @property
    def _final_estimator(self):
        return self.steps[-1][1]

    # Estimator interface

    def _pre_transform(self, X, y=None, **fit_params):
        fit_params_steps = dict((step, {}) for step, _ in self.steps)
        for pname, pval in six.iteritems(fit_params):
            step, param = pname.split('__', 1)
            fit_params_steps[step][param] = pval
        Xt = X
        for name, transform in self.steps[:-1]:
            if hasattr(transform, "fit_transform"):
                Xt = transform.fit_transform(Xt, y, **fit_params_steps[name])
            else:
                Xt = transform.fit(Xt, y, **fit_params_steps[name]) \
                              .transform(Xt)
        return Xt, fit_params_steps[self.steps[-1][0]]

    def fit(self, X, y=None, **fit_params):
        """Fit all the transforms one after the other and transform the
        data, then fit the transformed data using the final estimator.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of the
            pipeline.
        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        self.steps[-1][-1].fit(Xt, y, **fit_params)
        return self

    def fit_transform(self, X, y=None, **fit_params):
        """Fit all the transforms one after the other and transform the
        data, then use fit_transform on transformed data using the final
        estimator.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of the
            pipeline.

        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        if hasattr(self.steps[-1][-1], 'fit_transform'):
            return self.steps[-1][-1].fit_transform(Xt, y, **fit_params)
        else:
            return self.steps[-1][-1].fit(Xt, y, **fit_params).transform(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict(self, X):
        """Applies transforms to the data, and the predict method of the
        final estimator. Valid only if the final estimator implements
        predict.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def fit_predict(self, X, y=None, **fit_params):
        """Applies fit_predict of last step in pipeline after transforms.

        Applies fit_transforms of a pipeline to the data, followed by the
        fit_predict method of the final estimator in the pipeline. Valid
        only if the final estimator implements fit_predict.

        Parameters
        ----------
        X : iterable
            Training data. Must fulfill input requirements of first step of
            the pipeline.
        y : iterable, default=None
            Training targets. Must fulfill label requirements for all steps
            of the pipeline.
        """
        Xt, fit_params = self._pre_transform(X, y, **fit_params)
        return self.steps[-1][-1].fit_predict(Xt, y, **fit_params)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict_proba(self, X):
        """Applies transforms to the data, and the predict_proba method of the
        final estimator. Valid only if the final estimator implements
        predict_proba.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict_proba(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def decision_function(self, X):
        """Applies transforms to the data, and the decision_function method of
        the final estimator. Valid only if the final estimator implements
        decision_function.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].decision_function(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def predict_log_proba(self, X):
        """Applies transforms to the data, and the predict_log_proba method of
        the final estimator. Valid only if the final estimator implements
        predict_log_proba.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].predict_log_proba(Xt)

    @if_delegate_has_method(delegate='_final_estimator')
    def transform(self, X):
        """Applies transforms to the data, and the transform method of the
        final estimator. Valid only if the final estimator implements
        transform.

        Parameters
        ----------
        X : iterable
            Data to predict on. Must fulfill input requirements of first step of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps:
            Xt = transform.transform(Xt)
        return Xt

    @if_delegate_has_method(delegate='_final_estimator')
    def inverse_transform(self, X):
        """Applies inverse transform to the data.
        Starts with the last step of the pipeline and applies ``inverse_transform`` in
        inverse order of the pipeline steps.
        Valid only if all steps of the pipeline implement inverse_transform.

        Parameters
        ----------
        X : iterable
            Data to inverse transform. Must fulfill output requirements of the
            last step of the pipeline.
        """
        if X.ndim == 1:
            X = X[None, :]
        Xt = X
        for name, step in self.steps[::-1]:
            Xt = step.inverse_transform(Xt)
        return Xt

    @if_delegate_has_method(delegate='_final_estimator')
    def score(self, X, y=None):
        """Applies transforms to the data, and the score method of the
        final estimator. Valid only if the final estimator implements
        score.

        Parameters
        ----------
        X : iterable
            Data to score. Must fulfill input requirements of first step of the
            pipeline.

        y : iterable, default=None
            Targets used for scoring. Must fulfill label requirements for all steps of
            the pipeline.
        """
        Xt = X
        for name, transform in self.steps[:-1]:
            Xt = transform.transform(Xt)
        return self.steps[-1][-1].score(Xt, y)

    @property
    def classes_(self):
        return self.steps[-1][-1].classes_

    @property
    def _pairwise(self):
        # check if first estimator expects pairwise input
        return getattr(self.steps[0][1], '_pairwise', False)

def _name_estimators(estimators):
    """Generate names for estimators."""

    names = [type(estimator).__name__.lower() for estimator in estimators]
    namecount = defaultdict(int)
    for est, name in zip(estimators, names):
        namecount[name] += 1

    for k, v in list(six.iteritems(namecount)):
        if v == 1:
            del namecount[k]

    for i in reversed(range(len(estimators))):
        name = names[i]
        if name in namecount:
            names[i] += "-%d" % namecount[name]
            namecount[name] -= 1

    return list(zip(names, estimators))

def make_pipeline(*steps):
    """Construct a Pipeline from the given estimators.

    This is a shorthand for the Pipeline constructor; it does not require, and
    does not permit, naming the estimators. Instead, they will be given names
    automatically based on their types.

    Examples
    --------
    >>> from sklearn.naive_bayes import GaussianNB
    >>> from sklearn.preprocessing import StandardScaler
    >>> make_pipeline(StandardScaler(), GaussianNB())    # doctest: +NORMALIZE_WHITESPACE
    Pipeline(steps=[('standardscaler',
                     StandardScaler(copy=True, with_mean=True, with_std=True)),
                    ('gaussiannb', GaussianNB())])

    Returns
    -------
    p : Pipeline
    """
    return Pipeline(_name_estimators(steps))

def _fit_one_transformer(transformer, X, y):
    return transformer.fit(X, y)

def _transform_one(transformer, name, X, transformer_weights):
    if transformer_weights is not None and name in transformer_weights:
        # if we have a weight for this transformer, muliply output
        return transformer.transform(X) * transformer_weights[name]
    return transformer.transform(X)

def _fit_transform_one(transformer, name, X, y, transformer_weights,
                       **fit_params):
    if transformer_weights is not None and name in transformer_weights:
        # if we have a weight for this transformer, muliply output
        if hasattr(transformer, 'fit_transform'):
            X_transformed = transformer.fit_transform(X, y, **fit_params)
            return X_transformed * transformer_weights[name], transformer
        else:
            X_transformed = transformer.fit(X, y, **fit_params).transform(X)
            return X_transformed * transformer_weights[name], transformer
    if hasattr(transformer, 'fit_transform'):
        X_transformed = transformer.fit_transform(X, y, **fit_params)
        return X_transformed, transformer
    else:
        X_transformed = transformer.fit(X, y, **fit_params).transform(X)
        return X_transformed, transformer

class FeatureUnion(BaseEstimator, TransformerMixin):
    """Concatenates results of multiple transformer objects.

    This estimator applies a list of transformer objects in parallel to the
    input data, then concatenates the results. This is useful to combine
    several feature extraction mechanisms into a single transformer.

    Read more in the :ref:`User Guide <feature_union>`.

    Parameters
    ----------
    transformer_list: list of (string, transformer) tuples
        List of transformer objects to be applied to the data. The first
        half of each tuple is the name of the transformer.

    n_jobs: int, optional
        Number of jobs to run in parallel (default 1).

    transformer_weights: dict, optional
        Multiplicative weights for features per transformer.
        Keys are transformer names, values the weights.

    """
    def __init__(self, transformer_list, n_jobs=1, transformer_weights=None):
        self.transformer_list = transformer_list
        self.n_jobs = n_jobs
        self.transformer_weights = transformer_weights

    def get_feature_names(self):
        """Get feature names from all transformers.

        Returns
        -------
        feature_names : list of strings
            Names of the features produced by transform.
        """
        feature_names = []
        for name, trans in self.transformer_list:
            if not hasattr(trans, 'get_feature_names'):
                raise AttributeError("Transformer %s does not provide"
                                     " get_feature_names." % str(name))
            feature_names.extend([name + "__" + f for f in
                                  trans.get_feature_names()])
        return feature_names

    def fit(self, X, y=None):
        """Fit all transformers using X.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data, used to fit transformers.
        """
        transformers = Parallel(n_jobs=self.n_jobs)(
            delayed(_fit_one_transformer)(trans, X, y)
            for name, trans in self.transformer_list)
        self._update_transformer_list(transformers)
        return self

    def fit_transform(self, X, y=None, **fit_params):
        """Fit all transformers using X, transform the data and concatenate
        results.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data to be transformed.

        Returns
        -------
        X_t : array-like or sparse matrix, shape (n_samples, sum_n_components)
            hstack of results of transformers. sum_n_components is the
            sum of n_components (output dimension) over transformers.
        """
        result = Parallel(n_jobs=self.n_jobs)(
            delayed(_fit_transform_one)(trans, name, X, y,
                                        self.transformer_weights, **fit_params)
            for name, trans in self.transformer_list)

        Xs, transformers = zip(*result)
        self._update_transformer_list(transformers)
        if any(sparse.issparse(f) for f in Xs):
            Xs = sparse.hstack(Xs).tocsr()
        else:
            Xs = np.hstack(Xs)
        return Xs

    def transform(self, X):
        """Transform X separately by each transformer, concatenate results.

        Parameters
        ----------
        X : array-like or sparse matrix, shape (n_samples, n_features)
            Input data to be transformed.

        Returns
        -------
        X_t : array-like or sparse matrix, shape (n_samples, sum_n_components)
            hstack of results of transformers. sum_n_components is the
            sum of n_components (output dimension) over transformers.
        """
        Xs = Parallel(n_jobs=self.n_jobs)(
            delayed(_transform_one)(trans, name, X, self.transformer_weights)
            for name, trans in self.transformer_list)
        if any(sparse.issparse(f) for f in Xs):
            Xs = sparse.hstack(Xs).tocsr()
        else:
            Xs = np.hstack(Xs)
        return Xs

    def get_params(self, deep=True):
        if not deep:
            return super(FeatureUnion, self).get_params(deep=False)
        else:
            out = dict(self.transformer_list)
            for name, trans in self.transformer_list:
                for key, value in iteritems(trans.get_params(deep=True)):
                    out['%s__%s' % (name, key)] = value
            out.update(super(FeatureUnion, self).get_params(deep=False))
            return out

    def _update_transformer_list(self, transformers):
        self.transformer_list[:] = [
            (name, new)
            for ((name, old), new) in zip(self.transformer_list, transformers)
        ]

# XXX it would be nice to have a keyword-only n_jobs argument to this function,
# but that's not allowed in Python 2.x.
def make_union(*transformers):
    """Construct a FeatureUnion from the given transformers.

    This is a shorthand for the FeatureUnion constructor; it does not require,
    and does not permit, naming the transformers. Instead, they will be given
    names automatically based on their types. It also does not allow weighting.

    Examples
    --------
    >>> from sklearn.decomposition import PCA, TruncatedSVD
    >>> make_union(PCA(), TruncatedSVD())    # doctest: +NORMALIZE_WHITESPACE
    FeatureUnion(n_jobs=1,
                 transformer_list=[('pca', PCA(copy=True, n_components=None,
                                               whiten=False)),
                                   ('truncatedsvd',
                                    TruncatedSVD(algorithm='randomized',
                                                 n_components=2, n_iter=5,
                                                 random_state=None, tol=0.0))],
                 transformer_weights=None)

    Returns
    -------
    f : FeatureUnion
    """
    return FeatureUnion(_name_estimators(transformers))


#   tftools.py: Utilities for optimizing transfer function excitation signals
#   Copyright (C) 2013  Larry Price
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import division

import numpy as np
import scipy.signal as sig
from matplotlib.mlab import csd, psd
import sys

def tukeywin(m, a=0.5):
	'''
	Produces a tukey window

	a = overlap parameter between 0 returns a square window and 1 returns a Hann(ing) window
	m = number of points in the window

	see, e.g., https://en.wikipedia.org/wiki/Window_function#Tukey_window
	'''

	if a <= 0:
		return np.ones(m)
	elif a >= 1:
		return np.hanning(m)

	x = np.linspace(0, 1, m)
	w = np.ones_like(x)
	w[x < a/2] = (1 + np.cos(2*np.pi/a * (x[x < a/2] - a/2) )) / 2
	w[x >= (1 - a/2)] = (1 + np.cos(2*np.pi/a * (x[x >= (1 - a/2)] - 1 + a/2))) / 2

	return w

def getTF(exc,resp,tmeas,Fs,Nfft,padto=None):
	"""
	compute transfer function along with coherence and SNR. uses PSD/CSD method with 50% overlapping windows
	returns f, TF, Coherence, SNR

	exc = excitation signal
	resp = system response
	tmeas = duration of mesurement in seconds
	Fs = sample rate (Hz)
	Nfft = number of data points to be used in each block for fft
	padto = pad to this many points
	"""

	N = 1.89 * tmeas / (Nfft / Fs)

	Sx, f = psd(exc,NFFT=Nfft,Fs=Fs,noverlap=int(Nfft/2),pad_to=padto)
	Sy = psd(resp,NFFT=Nfft,Fs=Fs,noverlap=int(Nfft/2),pad_to=padto)[0]
	Sxy = csd(exc,resp,NFFT=Nfft,Fs=Fs,noverlap=int(Nfft/2),pad_to=padto)[0]
	Cxy = (Sxy * np.conj(Sxy)) / (Sx * Sy)
	snr = np.sqrt(Cxy * 2 * N / (1 - Cxy) )

	return f, Sxy / Sx, Cxy, snr

def fishersfzpk(ltisys,w,Sn):
	"""
	create the single-frequency fisher matrix for a transfer function in zpk form, i.e.
	           __
			   ||_i (w - z_i)

	H(w) = k --------------------
	           __
			   ||_i (w - p_i)

	*** the excitation signal is assumed to be a sum of sines ***
	*** the denominator must be monic ***

	arguments:

	ltisys = a scipy.signal.lti instance of the transfer function
	w = (angular) frequencies of interest (in rad/s)
	Sn = PSD of the noise as an array (same length as w)

	returns:
		an NxN numpy array with each value being an array of len(w) (the collection of all single
		frequency Fisher matrices at frequencies w)
	"""
	###FIXME: add some basic error handling

	#tf is in terms of iw, not w
	s = 1j*w

	#create the transfer function
	#use the lower case w because scipy puts the i in the transfer function for us
	tf = ltisys.freqresp(w)[1]

	#do the magnitude squared here once and for all
	tfmagsq = tf * np.conj(tf)

	#get the number of parameters
	if ltisys.gain == 1:
		N = len(ltisys.zeros) + len(ltisys.poles)
	else:
		N = len(ltisys.zeros) + len(ltisys.poles) + 1

	#take all the derivatives
	Dz = np.zeros([len(ltisys.zeros),len(s)],dtype=np.complex128)
	Dp = np.zeros([len(ltisys.poles),len(s)],dtype=np.complex128)

	for i,z in enumerate(ltisys.zeros):
		Dz[i] = -1 / (s - z)
	for i,p in enumerate(ltisys.poles):
		Dp[i] = 1 / (s - p)

	#check for unity gain and the absence of zeros
	if ltisys.gain == 1 and ltisys.zeros.size:
		deez = np.vstack((Dz,Dp))
	elif ltisys.gain == 1 and not ltisys.zeros.size:
		deez = Dp
	elif ltisys.gain != 1 and ltisys.zeros.size:
		deez = np.vstack((np.vstack((Dz,Dp)), 1/ltisys.gain[0] * np.ones(len(s))))
	else:
		deez = np.vstack((Dp,1/ltisys.gain[0] * np.ones(len(s))))

	#put it together to make the fisher matrix
	fisher = np.zeros([N,N,len(w)],dtype=np.float64)
	for i in range(N):
		for j in range(N):
			fisher[i][j] = 0.5 * tfmagsq * np.real(np.conj(deez[i])*deez[j]) / Sn

	#all done
	return fisher

def fishersfab(ltisys,w,Sn):
	"""
	create the single-frequency fisher matrix for a transfer function as a rational function, i.e.

			  Sum_0^N b_i s^i
	H(w) =  --------------------
			1 + Sum_1^M a_i s^i

	*** the excitation signal is assumed to be a sum of sines ***
	*** the denominator must be monic (it's enforced, so no worries)***

	arguments:

	ltisys = instance of scipy.signal.lti
	w = frequencies of interest
	Sn = PSD of the noise as an array (same length as w)

	returns:
		an NxN numpy array with each value being an array of len(w) (the collection of all single
		frequency Fisher matrices at frequencies w)
		NB: you have to take the transpose of the result if you want to, say compute the determinant via numpy.linalg.det
	"""
	###FIXME: add some basic error handling

	#tf is in terms of iw, not w
	s = 1j*w

	#get the tf in the right form
	a,b = lti2ab(ltisys)

	#create the numerator and denominator of the tf
	if b.size:
		numer = np.sum(np.array([ b[i] * s**i for i in range(len(b))]),axis=0)
	else: #i don't think this is even possible unless the tf is just a pass-through
		numer = np.ones(len(s))

	denom = np.sum(np.array([ a[i] * s**(i+1) for i in range(len(a))]),axis=0) + np.ones(len(s))

	#number of parameters
	N = len(a) + len(b)

	#take all the derivatives
	deez = np.zeros([N,len(w)],dtype=np.complex128)

	for i in range(N):
		#derivative wrt denominator
		#funky numbering because denom is monic (demonic?)
		if i < len(a):
			deez[i] = - s**(i+1) * numer / denom**2
		#derivative wrt numerator
		else:
			deez[i] = s**(i-len(a)) / denom

	#put it together to make the fisher matrix
	fisher = np.zeros([N,N,len(w)],dtype=np.float64)
	for i in range(N):
		for j in range(N):
			fisher[i][j] = 0.5 * np.real(np.conj(deez[i])*deez[j]) / Sn

	#all done
	return fisher

def fishersf(ltisys,w,Sn,usezpk=False):
	"""
	convenience function to select between zpk (default) and ab form for computing the fisher matrix
	"""
	if usezpk is True:
		return fishersfzpk(ltisys,w,Sn)
	else:
		return fishersfab(ltisys,w,Sn)

def fisherdesign(fmsf,Sx):
	"""
	compute the fisher matrix associated with the design Sx
	uses the Sx and the single frequency fisher matrix
	"""

	return np.sum(fmsf*Sx,axis=2)

def dispersion(fmdesign,fmsf):
	"""
	compute the dispersion from the single frequency and design fisher matrices
	"""

	return np.trace(np.dot(np.linalg.inv(fmdesign),fmsf))

def lti2ab(ltisys):
	"""
	convenience function to convert scipy.signal.lti instance to a,b suitable for fisher matrix calculation
	ltisys is an instance of scipy.signal.lti

	returns a,b
	"""

	b = ltisys.num
	a = ltisys.den

	#fancy array slicing to reverse the order (::-1) and remove the first element of a (1:)
	return a[::-1][1:] / a[-1], b[::-1] / a[-1]

def findfreqs(ltisys,Sn,w,nfreqs=None,usezpk=False):
	"""
	find best frequencies for optimal design (brute force method)

	arguments:

	ltisys = instance of scipy.signal.lti
	w = (angular) frequencies of interest
	nfreqs = # of frequencies to return.  default is 3 x #parameters
	usezpk = boolean for indicating form of the transfer function

	returns:

	wopt = array of optimal frequencies to use
	fisherf = single-frequency fisher matrix evaluated at wopt (basically input for design optimization)
	"""
	#get the number of parameters and put the transfer function in the right form
	if usezpk is True:
		#number of parameters for zpk representation
		if ltisys.gain == 1:
			nparm = len(ltisys.zeros) + len(ltisys.poles)
		else:
			nparm = len(ltisys.zeros) + len(ltisys.poles) + 1
	else:
		#using ab form
		a,b = lti2ab(ltisys)

		#number of parameters
		nparm = len(a) + len(b)

	#set the number of frequencies
	if nfreqs is None:
		nfreqs = 3 * nparm
	if nfreqs < 2 * nparm:
		raise ValueError('Must specify an nfreqs at least twice as large as the number of parameters!')
		sys.exit(0)

	fmsf = fishersf(ltisys,w,Sn,usezpk=usezpk)
	thesefreqs = np.sort(np.argsort(np.linalg.det(fmsf.T))[-nfreqs:])

	return w[thesefreqs], fmsf.T[thesefreqs].T

def optdesign(ltisys,w,usezpk=False,fmsf=None,Sn=None,tol=None,maxit=10000):
	"""
	compute the optimal design, Sx

	arguments:

	ltisys = instance of scipy.signal.lti
	w = the frequencies to optimize over
	tol = if max(dispersion - nparam) < tol, then iteration ceases.  if tol isn't specified then iteration continues until maxit
	maxit = maximum number of iterations to perform

	returns a tuple containing:

	Sx = optimal design as a numpy array
	max(dispersion - nparam)
	"""
	#FIXME: add some error handling
	if fmsf is None and Sn is None:
		raise ValueError('Must specify Sn to compute Fisher!')
		sys.exit(1)

	#get the number of parameters and put the transfer function in the right form
	if usezpk is True:
		#number of parameters for zpk representation
		if ltisys.gain == 1:
			nparm = len(ltisys.zeros) + len(ltisys.poles)
		else:
			nparm = len(ltisys.zeros) + len(ltisys.poles) + 1
	else:
		#using ab form
		a,b = lti2ab(ltisys)

		#number of parameters
		nparm = len(a) + len(b)

	#compute the single frequency fisher matrix
	if fmsf is None:
		fmsf = fishersf(ltisys,w,Sn,usezpk=usezpk)

	#initial design
	#normalized to one with all the power evenly distributed
	#don't worry about phases for now...FIXME: optimize phases
	Sx = np.ones(len(w)) / len (w)

	#compute the dispersion
	disp = dispersion(fisherdesign(fmsf,Sx),fmsf)

	for i in range(maxit):
		if tol is not None:
			if np.max(disp - nparm) < tol:
				break
		else:
			Sx *= (disp / nparm)
			disp = dispersion(fisherdesign(fmsf,Sx),fmsf)

	return Sx, np.max(disp - nparm)


"""
This example shows how to use a path patch to draw a bunch of
rectangles.  The technique of using lots of Rectangle instances, or
the faster method of using PolyCollections, were implemented before we
had proper paths with moveto/lineto, closepoly etc in mpl.  Now that
we have them, we can draw collections of regularly shaped objects with
homogeous properties more efficiently with a PathCollection.  This
example makes a histogram -- its more work to set up the vertex arrays
at the outset, but it should be much faster for large numbers of
objects
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.path as path

fig = plt.figure()
ax = fig.add_subplot(111)

# histogram our data with numpy
data = np.random.randn(1000)
n, bins = np.histogram(data, 50)

# get the corners of the rectangles for the histogram
left = np.array(bins[:-1])
right = np.array(bins[1:])
bottom = np.zeros(len(left))
top = bottom + n

# we need a (numrects x numsides x 2) numpy array for the path helper
# function to build a compound path
XY = np.array([[left,left,right,right], [bottom,top,top,bottom]]).T

# get the Path object
barpath = path.Path.make_compound_path_from_polys(XY)

# make a patch out of it
patch = patches.PathPatch(barpath, facecolor='blue', edgecolor='gray', alpha=0.8)
ax.add_patch(patch)

# update the view limits
ax.set_xlim(left[0], right[-1])
ax.set_ylim(bottom.min(), top.max())

plt.show()


from setuptools import setup, find_packages

version = '0.1.16'

setup(
    name='padua',
    version=version,
    url='http://github.com/mfitzp/padua',
    author='Martin Fitzpatrick',
    author_email='martin.fitzpatrick@gmail.com',
    description='A Python interface for Proteomic Data Analysis, working with MaxQuant & Perseus outputs',
    license='MIT',
    packages=find_packages(),
    include_package_data=True,
    classifiers=[
        'Development Status :: 2 - Pre-Alpha',
        'Intended Audience :: Developers',
        'Topic :: Desktop Environment',
        'Topic :: Software Development :: Build Tools',
        'Topic :: Software Development :: Widget Sets',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3.4'
    ],
    install_requires=[
        'numpy',
        'scipy',
        'matplotlib',
        'pandas',
        'statsmodels',
        'matplotlib-venn',
        'scikit-learn',
        'requests',
        'requests_toolbelt',
        'adjustText'
    ]
)


"""
Agglomerative clustering with different metrics
===============================================

Demonstrates the effect of different metrics on the hierarchical clustering.

The example is engineered to show the effect of the choice of different
metrics. It is applied to waveforms, which can be seen as
high-dimensional vector. Indeed, the difference between metrics is
usually more pronounced in high dimension (in particular for euclidean
and cityblock).

We generate data from three groups of waveforms. Two of the waveforms
(waveform 1 and waveform 2) are proportional one to the other. The cosine
distance is invariant to a scaling of the data, as a result, it cannot
distinguish these two waveforms. Thus even with no noise, clustering
using this distance will not separate out waveform 1 and 2.

We add observation noise to these waveforms. We generate very sparse
noise: only 6% of the time points contain noise. As a result, the
l1 norm of this noise (ie "cityblock" distance) is much smaller than it's
l2 norm ("euclidean" distance). This can be seen on the inter-class
distance matrices: the values on the diagonal, that characterize the
spread of the class, are much bigger for the Euclidean distance than for
the cityblock distance.

When we apply clustering to the data, we find that the clustering
reflects what was in the distance matrices. Indeed, for the Euclidean
distance, the classes are ill-separated because of the noise, and thus
the clustering does not separate the waveforms. For the cityblock
distance, the separation is good and the waveform classes are recovered.
Finally, the cosine distance does not separate at all waveform 1 and 2,
thus the clustering puts them in the same cluster.
"""
# Author: Gael Varoquaux
# License: BSD 3-Clause or CC-0

import matplotlib.pyplot as plt
import numpy as np

from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics import pairwise_distances

np.random.seed(0)

# Generate waveform data
n_features = 2000
t = np.pi * np.linspace(0, 1, n_features)

def sqr(x):
    return np.sign(np.cos(x))

X = list()
y = list()
for i, (phi, a) in enumerate([(.5, .15), (.5, .6), (.3, .2)]):
    for _ in range(30):
        phase_noise = .01 * np.random.normal()
        amplitude_noise = .04 * np.random.normal()
        additional_noise = 1 - 2 * np.random.rand(n_features)
        # Make the noise sparse
        additional_noise[np.abs(additional_noise) < .997] = 0

        X.append(12 * ((a + amplitude_noise)
                 * (sqr(6 * (t + phi + phase_noise)))
                 + additional_noise))
        y.append(i)

X = np.array(X)
y = np.array(y)

n_clusters = 3

labels = ('Waveform 1', 'Waveform 2', 'Waveform 3')

# Plot the ground-truth labelling
plt.figure()
plt.axes([0, 0, 1, 1])
for l, c, n in zip(range(n_clusters), 'rgb',
                   labels):
    lines = plt.plot(X[y == l].T, c=c, alpha=.5)
    lines[0].set_label(n)

plt.legend(loc='best')

plt.axis('tight')
plt.axis('off')
plt.suptitle("Ground truth", size=20)

# Plot the distances
for index, metric in enumerate(["cosine", "euclidean", "cityblock"]):
    avg_dist = np.zeros((n_clusters, n_clusters))
    plt.figure(figsize=(5, 4.5))
    for i in range(n_clusters):
        for j in range(n_clusters):
            avg_dist[i, j] = pairwise_distances(X[y == i], X[y == j],
                                                metric=metric).mean()
    avg_dist /= avg_dist.max()
    for i in range(n_clusters):
        for j in range(n_clusters):
            plt.text(i, j, '%5.3f' % avg_dist[i, j],
                     verticalalignment='center',
                     horizontalalignment='center')

    plt.imshow(avg_dist, interpolation='nearest', cmap=plt.cm.gnuplot2,
               vmin=0)
    plt.xticks(range(n_clusters), labels, rotation=45)
    plt.yticks(range(n_clusters), labels)
    plt.colorbar()
    plt.suptitle("Interclass %s distances" % metric, size=18)
    plt.tight_layout()

# Plot clustering results
for index, metric in enumerate(["cosine", "euclidean", "cityblock"]):
    model = AgglomerativeClustering(n_clusters=n_clusters,
                                    linkage="average", affinity=metric)
    model.fit(X)
    plt.figure()
    plt.axes([0, 0, 1, 1])
    for l, c in zip(np.arange(model.n_clusters), 'rgbk'):
        plt.plot(X[model.labels_ == l].T, c=c, alpha=.5)
    plt.axis('tight')
    plt.axis('off')
    plt.suptitle("AgglomerativeClustering(affinity=%s)" % metric, size=20)

plt.show()


from statsmodels.tsa import arima_model
import numpy as np
from pug.invest import util

y = util.simulate(poly=100, sinusoids=(10, 100, -20)).values

hr = np.arange(365*96)*.25
t = hr * 3600
sinusoids = [
    np.random.normal(0.0, 0.1, 365*96)+10 + 3*np.sin(hr*2*np.pi/96/.25),
    np.random.normal(0.0, 0.1, 365*96)+15 + 3*np.sin(hr*2*np.pi/96/.25) + 3*np.cos(t*2*np.pi/96./.25/365.),
    np.random.normal(0.0, 1.0, 365*96)+15 + 3*np.sin(hr*2*np.pi/96/.25) + 3*np.cos(t*2*np.pi/96./.25/365.)+np.random.normal(0.0,1e-5,365*96).cumsum()]
arma20 = arima_model.ARMA(y, (2,0)).fit()
y2 = arma.predict(start=10*96, end=12*96)
y1 = y[10*96-1:12*96]
plt.plot(t[10*96-1:12*96],zip(*[y1,y2]))
plt.show()
y2 = arma30.predict(start=10*96, end=12*96)
plt.plot(t[10*96-1:12*96],zip(*[y1,y2]))
plt.show()
arma30.resid.plot()
plt.plot(arma30.resid)
plt.show()
plt.plot(arma30.resid/y2)
plt.plot(arma30.resid/y)
plt.show()
plt.plot(arma30.resid/y)
plt.show()
arma30 = arima_model.ARMA(y[:-96*30], (2,0)).fit()
y1 = y[-32*96:]
y2 = arma30.predict(start=N-32*96, end=N-28*96)
N=len(y)
y2 = arma30.predict(start=N-32*96, end=N-28*96)
plt.plot(t[-32*96-1:-28*96],zip(*[y1,y2]))
plt.show()
plt.plot(t[-32*96-1:-28*96],zip(*[y1,y2]))
plt.show()
N
arma30 = arima_model.ARMA(y[:-96*30], (3,0)).fit()
N_predict=len(y[:-96*30])
y_predict=y[:-96*30]
y2 = arma30.predict(start=N_predict,end=N_predict+96)
y1 = y[N_predict:N_predict+96]
y1-y2
y2 = arma30.predict(start=N_predict,end=N_predict+95)
plt.plot(zip(*[y1,y2]))
plt.plot(zip(*[y1,y2]))
plt.show()
arma41 = arima_model.ARMA(y_train, (4,1)).fit()
y_train=y[:-96*30]
arma41 = arima_model.ARMA(y_train, (4,1)).fit()

arma296 = arima_model.ARMA(y_train, (2,96)).fit()
arma296 = arima_model.ARMA(y_train.diff(), (2,96)).fit()
arma296 = arima_model.ARMA(pd.Series(y_train).diff(), (2,96)).fit()
import pandas as pd
y_diff = pd.Series(y).diff().values()
y_diff = pd.Series(y).diff().values
y_train=y_diff[:-96*30]
arma296 = arima_model.ARMA(y_train, (2,96)).fit()
arma296 = arima_model.ARMA(y_train, (2,0)).fit()
arma296 = arima_model.ARMA(y_train[1:], (2,0)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (2,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*7:], (2,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*2:], (2,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*3:], (2,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*4:], (2,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (2,48)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (2,24)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (0,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (1,96)).fit()
arma296 = arima_model.ARMA(y_train[-96*14:], (1,96)).fit(meth='mle')
arma296 = arima_model.ARMA(y_train[-96*14:], (1,96)).fit(meth='css')
arma296 = arima_model.ARMA(np.diff(y_train[-96*14:]).dropna(), (1,96)).fit(meth='css')
arma296 = arima_model.ARMA(np.diff(y_train[-96*14:])[1:], (2,96)).fit(meth='css')
arma296 = arima_model.ARMA(np.diff(y_train[-96*14:])[1:], (2,96)).fit(meth='mle')
arma296 = arima_model.ARMA(np.diff(y_train[-96*14:])[1:], (2,96))
arma296.fit(trend='c',solver='bfgs')
arma296.fit(trend='c',solver='bfgs',transparams=True)
arma296.fit(trend='c',solver='bfgs',transparams=False)
arma296._fit_start_params
arma296._fit_start_params()
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False)
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=True)
q = np.zeros(96)
q[0] = 1
q[-1]=1
q[-1]=.5
q[0] = .1
q[-1]=.9
p=[10, 1.2, -.2]
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=True,startparams=[p,q])
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=True,start_params=[p,q])
np.log
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False,start_params=[p,q])
p=np.array([10, 1.2, -.2])
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False,start_params=[p,q])
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False,start_params=np.array([p,q]))
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False,start_params=q)
q.shape
q = np.zeros(93)
q[-1]=.9
q[0]=.1
arma296.fit(meth='css-mle',trend='c',solver='bfgs',transparams=False,start_params=q)
arma296.fit(trend='c',solver='bfgs',transparams=False,start_params=q)
arma296.fit(trend='c',transparams=False,start_params=q)
arma296.fit(transparams=False,start_params=q)
len(q)
p=np.array([10, 1.2, -.2])
q = np.zeros(99)
q[0]=.1
q[0]=10
q[1]=1
q[2]=-.2
q[-1]=.95
arma296.fit(transparams=False,start_params=q)


"""
====================================================================
Linear and Quadratic Discriminant Analysis with covariance ellipsoid
====================================================================

This example plots the covariance ellipsoids of each class and
decision boundary learned by LDA and QDA. The ellipsoids display
the double standard deviation for each class. With LDA, the
standard deviation is the same for all the classes, while each
class has its own standard deviation with QDA.
"""
print(__doc__)

from scipy import linalg
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib import colors

from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis

###############################################################################
# colormap
cmap = colors.LinearSegmentedColormap(
    'red_blue_classes',
    {'red': [(0, 1, 1), (1, 0.7, 0.7)],
     'green': [(0, 0.7, 0.7), (1, 0.7, 0.7)],
     'blue': [(0, 0.7, 0.7), (1, 1, 1)]})
plt.cm.register_cmap(cmap=cmap)

###############################################################################
# generate datasets
def dataset_fixed_cov():
    '''Generate 2 Gaussians samples with the same covariance matrix'''
    n, dim = 300, 2
    np.random.seed(0)
    C = np.array([[0., -0.23], [0.83, .23]])
    X = np.r_[np.dot(np.random.randn(n, dim), C),
              np.dot(np.random.randn(n, dim), C) + np.array([1, 1])]
    y = np.hstack((np.zeros(n), np.ones(n)))
    return X, y

def dataset_cov():
    '''Generate 2 Gaussians samples with different covariance matrices'''
    n, dim = 300, 2
    np.random.seed(0)
    C = np.array([[0., -1.], [2.5, .7]]) * 2.
    X = np.r_[np.dot(np.random.randn(n, dim), C),
              np.dot(np.random.randn(n, dim), C.T) + np.array([1, 4])]
    y = np.hstack((np.zeros(n), np.ones(n)))
    return X, y

###############################################################################
# plot functions
def plot_data(lda, X, y, y_pred, fig_index):
    splot = plt.subplot(2, 2, fig_index)
    if fig_index == 1:
        plt.title('Linear Discriminant Analysis')
        plt.ylabel('Data with fixed covariance')
    elif fig_index == 2:
        plt.title('Quadratic Discriminant Analysis')
    elif fig_index == 3:
        plt.ylabel('Data with varying covariances')

    tp = (y == y_pred)  # True Positive
    tp0, tp1 = tp[y == 0], tp[y == 1]
    X0, X1 = X[y == 0], X[y == 1]
    X0_tp, X0_fp = X0[tp0], X0[~tp0]
    X1_tp, X1_fp = X1[tp1], X1[~tp1]

    alpha = 0.5

    # class 0: dots
    plt.plot(X0_tp[:, 0], X0_tp[:, 1], 'o', alpha=alpha,
             color='red')
    plt.plot(X0_fp[:, 0], X0_fp[:, 1], '*', alpha=alpha,
             color='#990000')  # dark red

    # class 1: dots
    plt.plot(X1_tp[:, 0], X1_tp[:, 1], 'o', alpha=alpha,
             color='blue')
    plt.plot(X1_fp[:, 0], X1_fp[:, 1], '*', alpha=alpha,
             color='#000099')  # dark blue

    # class 0 and 1 : areas
    nx, ny = 200, 100
    x_min, x_max = plt.xlim()
    y_min, y_max = plt.ylim()
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, nx),
                         np.linspace(y_min, y_max, ny))
    Z = lda.predict_proba(np.c_[xx.ravel(), yy.ravel()])
    Z = Z[:, 1].reshape(xx.shape)
    plt.pcolormesh(xx, yy, Z, cmap='red_blue_classes',
                   norm=colors.Normalize(0., 1.))
    plt.contour(xx, yy, Z, [0.5], linewidths=2., colors='k')

    # means
    plt.plot(lda.means_[0][0], lda.means_[0][1],
             'o', color='black', markersize=10)
    plt.plot(lda.means_[1][0], lda.means_[1][1],
             'o', color='black', markersize=10)

    return splot

def plot_ellipse(splot, mean, cov, color):
    v, w = linalg.eigh(cov)
    u = w[0] / linalg.norm(w[0])
    angle = np.arctan(u[1] / u[0])
    angle = 180 * angle / np.pi  # convert to degrees
    # filled Gaussian at 2 standard deviation
    ell = mpl.patches.Ellipse(mean, 2 * v[0] ** 0.5, 2 * v[1] ** 0.5,
                              180 + angle, facecolor=color, edgecolor='yellow',
                              linewidth=2, zorder=2)
    ell.set_clip_box(splot.bbox)
    ell.set_alpha(0.5)
    splot.add_artist(ell)
    splot.set_xticks(())
    splot.set_yticks(())

def plot_lda_cov(lda, splot):
    plot_ellipse(splot, lda.means_[0], lda.covariance_, 'red')
    plot_ellipse(splot, lda.means_[1], lda.covariance_, 'blue')

def plot_qda_cov(qda, splot):
    plot_ellipse(splot, qda.means_[0], qda.covariances_[0], 'red')
    plot_ellipse(splot, qda.means_[1], qda.covariances_[1], 'blue')

###############################################################################
for i, (X, y) in enumerate([dataset_fixed_cov(), dataset_cov()]):
    # Linear Discriminant Analysis
    lda = LinearDiscriminantAnalysis(solver="svd", store_covariance=True)
    y_pred = lda.fit(X, y).predict(X)
    splot = plot_data(lda, X, y, y_pred, fig_index=2 * i + 1)
    plot_lda_cov(lda, splot)
    plt.axis('tight')

    # Quadratic Discriminant Analysis
    qda = QuadraticDiscriminantAnalysis(store_covariances=True)
    y_pred = qda.fit(X, y).predict(X)
    splot = plot_data(qda, X, y, y_pred, fig_index=2 * i + 2)
    plot_qda_cov(qda, splot)
    plt.axis('tight')
plt.suptitle('Linear Discriminant Analysis vs Quadratic Discriminant Analysis')
plt.show()


# general plot / graphics utility using matplotlib

from genaf.views.tools import *

from matplotlib import pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

import pandas
import io, base64

@roles( PUBLIC )
def index(request):

    # check

    if not request.GET.get('_method', None) in [ '_exec', '_dfexec' ]:

        pform, jscode = create_form( request )

        return render_to_response('genaf:templates/utils/index.mako',
            {   'title': 'Plotting Utility',
                'html': pform,
                'code': jscode,
            }, request = request )

    if request.GET.get('method') == '_dfexec':
        df = parse_df(request.GET.get('dfdata'))
    else:
        df = parse_textdata(request.GET.get('textdata'))
    plot_type = request.GET.get('plot_type')

    if plot_type == 'B':
        html, jscode = column_chart(df)
    elif plot_type == 'S':
        return error_page(request, 'Scatter plot not implemented yet')
    elif plot_type == 'P':
        html, jscode = pie_chart(df)

    return render_to_response('genaf:templates/utils/index.mako',
            {   'title': 'Plot',
                'html': html,
                'code': jscode,
            }, request = request )

def create_form(request):
    """ return html, jscode """

    pform = form(name='plotform', action='#')

    pform.add(
        fieldset(name='data')[
            input_textarea('textdata', label='Data'),
        ],
        fieldset(name='options')[
            input_select(name='plot_type', label='Plot type', value='B',
                    options = [ ('B', 'Bar (vertical) / column chart'),
                                ('S', 'Scatter x,y plot'),
                                ('P', 'Pie chart'),
                    ] ),
        ],
        fieldset()[ submit_bar('Create plot', '_exec')]
    )
    return (pform, '')

def parse_textdata(textdata):
    """ parse data, with the first line as header, and consecutive lines as data """
    header, content = textdata.split('\n', 1)
    columns = [ x.strip() for x in header.split('|') ]
    buff = io.StringIO(content)

    dataframe = pandas.read_table(buff, header=None, names = columns)

    return dataframe

def save_figure(canvas):

    figfile = io.BytesIO()
    canvas.print_figure(figfile)
    figfile.seek(0)
    figdata_png = figfile.getvalue()

    figdata_png = base64.b64encode(figdata_png).decode('ASCII')

    fig_html = literal('<img src="data:image/png;base64,%s" >' % figdata_png)

    return fig_html,''

def column_chart(df):
    """ creates column (vertical bar) chart """

    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(111)

    ax.bar(df.index, df.iloc[:,1], align='center')
    ax.set_xlabel(df.columns[0])
    ax.set_xticks(df.index)
    ax.set_xticklabels(df.iloc[:,0], rotation='vertical')
    ax.set_ylabel(df.columns[1])
    fig.tight_layout()

    return save_figure(canvas)

def pie_chart(df):

    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(111, aspect=1)

    ax.pie( df.iloc[:,1], labels = df.iloc[:,0], counterclock=False, startangle=90 )
    ax.set_xlabel(df.columns[0])
    fig.tight_layout()

    return save_figure(canvas)


# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------
"""
Provides two classes with the same signature for writing data out of NuPIC
models.
(This is a component of the One Hot Gym Anomaly Tutorial.)
"""
import csv
from collections import deque
from abc import ABCMeta, abstractmethod
from nupic.algorithms import anomaly_likelihood
# Try to import matplotlib, but we don't have to.
try:
  import matplotlib
  matplotlib.use('TKAgg')
  import matplotlib.pyplot as plt
  import matplotlib.gridspec as gridspec
  from matplotlib.dates import date2num, DateFormatter
except ImportError:
  pass

WINDOW = 300
HIGHLIGHT_ALPHA = 0.3
ANOMALY_HIGHLIGHT_COLOR = 'red'
WEEKEND_HIGHLIGHT_COLOR = 'yellow'
ANOMALY_THRESHOLD = 0.9

class NuPICOutput(object):

  __metaclass__ = ABCMeta

  def __init__(self, name):
    self.name = name
    self.anomalyLikelihoodHelper = anomaly_likelihood.AnomalyLikelihood()

  @abstractmethod
  def write(self, timestamp, value, predicted, anomalyScore):
    pass

  @abstractmethod
  def close(self):
    pass

class NuPICFileOutput(NuPICOutput):

  def __init__(self, *args, **kwargs):
    super(NuPICFileOutput, self).__init__(*args, **kwargs)
    self.outputFiles = []
    self.outputWriters = []
    self.lineCount = 0
    headerRow = [
      'timestamp', 'kw_energy_consumption', 'prediction',
      'anomaly_score', 'anomaly_likelihood'
    ]
    outputFileName = "%s_out.csv" % self.name
    print "Preparing to output %s data to %s" % (self.name, outputFileName)
    self.outputFile = open(outputFileName, "w")
    self.outputWriter = csv.writer(self.outputFile)
    self.outputWriter.writerow(headerRow)

  def write(self, timestamp, value, predicted, anomalyScore):
    if timestamp is not None:
      anomalyLikelihood = self.anomalyLikelihoodHelper.anomalyProbability(
        value, anomalyScore, timestamp
      )
      outputRow = [timestamp, value, predicted, anomalyScore, anomalyLikelihood]
      self.outputWriter.writerow(outputRow)
      self.lineCount += 1

  def close(self):
    self.outputFile.close()
    print "Done. Wrote %i data lines to %s." % (self.lineCount, self.name)

def extractWeekendHighlights(dates):
  weekendsOut = []
  weekendSearch = [5, 6]
  weekendStart = None
  for i, date in enumerate(dates):
    if date.weekday() in weekendSearch:
      if weekendStart is None:
        # Mark start of weekend
        weekendStart = i
    else:
      if weekendStart is not None:
        # Mark end of weekend
        weekendsOut.append((
          weekendStart, i, WEEKEND_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
        ))
        weekendStart = None

  # Cap it off if we're still in the middle of a weekend
  if weekendStart is not None:
    weekendsOut.append((
      weekendStart, len(dates)-1, WEEKEND_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
    ))

  return weekendsOut

def extractAnomalyIndices(anomalyLikelihood):
  anomaliesOut = []
  anomalyStart = None
  for i, likelihood in enumerate(anomalyLikelihood):
    if likelihood >= ANOMALY_THRESHOLD:
      if anomalyStart is None:
        # Mark start of anomaly
        anomalyStart = i
    else:
      if anomalyStart is not None:
        # Mark end of anomaly
        anomaliesOut.append((
          anomalyStart, i, ANOMALY_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
        ))
        anomalyStart = None

  # Cap it off if we're still in the middle of an anomaly
  if anomalyStart is not None:
    anomaliesOut.append((
      anomalyStart, len(anomalyLikelihood)-1,
      ANOMALY_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
    ))

  return anomaliesOut

class NuPICPlotOutput(NuPICOutput):

  def __init__(self, *args, **kwargs):
    super(NuPICPlotOutput, self).__init__(*args, **kwargs)
    # Turn matplotlib interactive mode on.
    plt.ion()
    self.dates = []
    self.convertedDates = []
    self.value = []
    self.allValues = []
    self.predicted = []
    self.anomalyScore = []
    self.anomalyLikelihood = []
    self.actualLine = None
    self.predictedLine = None
    self.anomalyScoreLine = None
    self.anomalyLikelihoodLine = None
    self.linesInitialized = False
    self._chartHighlights = []
    fig = plt.figure(figsize=(16, 10))
    gs = gridspec.GridSpec(2, 1, height_ratios=[3,  1])

    self._mainGraph = fig.add_subplot(gs[0, 0])
    plt.title(self.name)
    plt.ylabel('KW Energy Consumption')
    plt.xlabel('Date')

    self._anomalyGraph = fig.add_subplot(gs[1])

    plt.ylabel('Percentage')
    plt.xlabel('Date')

    # Maximizes window
    mng = plt.get_current_fig_manager()
    mng.resize(*mng.window.maxsize())

    plt.tight_layout()

  def initializeLines(self, timestamp):
    print "initializing %s" % self.name
    anomalyRange = (0.0, 1.0)
    self.dates = deque([timestamp] * WINDOW, maxlen=WINDOW)
    self.convertedDates = deque(
      [date2num(date) for date in self.dates], maxlen=WINDOW
    )
    self.value = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.predicted = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.anomalyScore = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.anomalyLikelihood = deque([0.0] * WINDOW, maxlen=WINDOW)

    actualPlot, = self._mainGraph.plot(self.dates, self.value)
    self.actualLine = actualPlot
    predictedPlot, = self._mainGraph.plot(self.dates, self.predicted)
    self.predictedLine = predictedPlot
    self._mainGraph.legend(tuple(['actual', 'predicted']), loc=3)

    anomalyScorePlot, = self._anomalyGraph.plot(
      self.dates, self.anomalyScore, 'm'
    )
    anomalyScorePlot.axes.set_ylim(anomalyRange)

    self.anomalyScoreLine = anomalyScorePlot
    anomalyLikelihoodPlot, = self._anomalyGraph.plot(
      self.dates, self.anomalyScore, 'r'
    )
    anomalyLikelihoodPlot.axes.set_ylim(anomalyRange)
    self.anomalyLikelihoodLine = anomalyLikelihoodPlot
    self._anomalyGraph.legend(
      tuple(['anomaly score', 'anomaly likelihood']), loc=3
    )

    dateFormatter = DateFormatter('%m/%d %H:%M')
    self._mainGraph.xaxis.set_major_formatter(dateFormatter)
    self._anomalyGraph.xaxis.set_major_formatter(dateFormatter)

    self._mainGraph.relim()
    self._mainGraph.autoscale_view(True, True, True)

    self.linesInitialized = True

  def highlightChart(self, highlights, chart):
    for highlight in highlights:
      # Each highlight contains [start-index, stop-index, color, alpha]
      self._chartHighlights.append(chart.axvspan(
        self.convertedDates[highlight[0]], self.convertedDates[highlight[1]],
        color=highlight[2], alpha=highlight[3]
      ))

  def write(self, timestamp, value, predicted, anomalyScore):

    # We need the first timestamp to initialize the lines at the right X value,
    # so do that check first.
    if not self.linesInitialized:
      self.initializeLines(timestamp)

    anomalyLikelihood = self.anomalyLikelihoodHelper.anomalyProbability(
      value, anomalyScore, timestamp
    )

    self.dates.append(timestamp)
    self.convertedDates.append(date2num(timestamp))
    self.value.append(value)
    self.allValues.append(value)
    self.predicted.append(predicted)
    self.anomalyScore.append(anomalyScore)
    self.anomalyLikelihood.append(anomalyLikelihood)

    # Update main chart data
    self.actualLine.set_xdata(self.convertedDates)
    self.actualLine.set_ydata(self.value)
    self.predictedLine.set_xdata(self.convertedDates)
    self.predictedLine.set_ydata(self.predicted)
    # Update anomaly chart data
    self.anomalyScoreLine.set_xdata(self.convertedDates)
    self.anomalyScoreLine.set_ydata(self.anomalyScore)
    self.anomalyLikelihoodLine.set_xdata(self.convertedDates)
    self.anomalyLikelihoodLine.set_ydata(self.anomalyLikelihood)

    # Remove previous highlighted regions
    for poly in self._chartHighlights:
      poly.remove()
    self._chartHighlights = []

    weekends = extractWeekendHighlights(self.dates)
    anomalies = extractAnomalyIndices(self.anomalyLikelihood)

    # Highlight weekends in main chart
    self.highlightChart(weekends, self._mainGraph)

    # Highlight anomalies in anomaly chart
    self.highlightChart(anomalies, self._anomalyGraph)

    maxValue = max(self.allValues)
    self._mainGraph.relim()
    self._mainGraph.axes.set_ylim(0, maxValue + (maxValue * 0.02))

    self._mainGraph.relim()
    self._mainGraph.autoscale_view(True, scaley=False)
    self._anomalyGraph.relim()
    self._anomalyGraph.autoscale_view(True, True, True)

    plt.draw()

  def close(self):
    plt.ioff()
    plt.show()

NuPICOutput.register(NuPICFileOutput)
NuPICOutput.register(NuPICPlotOutput)