# global imports
import numpy as np
from future.builtins import range

# local imports
import correlation_toolbox.helper as cthlp

'''
Documentation:

Correlation toolbox for AnalogSignals (binned data) of format

data = np.array([[t1,t2,...,tn], # unit1
                 [t1,t2,...,tn], # unit2
                 .
                 .
                 [t1,t2,...,tn]]) # unitN

Exception: compound_crossspec takes list of data, i.e. [data1,data2,...] as input.
'''


def mean(data, units=False, time=False):
    '''
    Compute mean of data

    **Args**:
       data: numpy.ndarray; 1st axis unit, 2nd axis time
       units: bool; average over units
       time: bool; average over time

    **Return**:
       if units=False and time=False: error, 
       if units=True: 1 dim numpy.ndarray; time series
       if time=True: 1 dim numpy.ndarray; series of unit means across time
       if units=True and time=True: float; unit and time mean

    **Examples**:
       >>> mean(np.array([[1,2,3],[4,5,6]]),units=True)
       Out[1]: np.array([2.5,3.5,4.5])

       >>> mean(np.array([[1,2,3],[4,5,6]]),time=True)
       Out[1]: np.array([2.,5.])

       >>> mean(np.array([[1,2,3],[4,5,6]]),units=True,time=True)
       Out[1]: 3.5

    '''

    assert(units is not False or time is not False)
    if units is True and time is False:
        return np.mean(data, axis=0)
    elif units is False and time is True:
        return np.mean(data, axis=1)
    elif units is True and time is True:
        return np.mean(data)


def compound_mean(data):
    '''
    Compute the mean of the compound/sum signal.
    data is first summed across units and averaged across time.

    **Args**:
       data: numpy.ndarray; 1st axis unit, 2nd axis time

    **Return**:
       float; time-averaged compound/sum signal
   
    **Examples**: 
       >>> compound_mean(np.array([[1,2,3],[4,5,6]]))
       Out[1]: 7.0

    '''

    return np.mean(np.sum(data, axis=0))


def variance(data, units=False, time=False):
    '''
    Compute the variance of data.

    **Args**:
       data: numpy.ndarray; 1st axis unit, 2nd axis time
       units: bool; variance across units
       time: bool; average over time

    **Return**:
       if units=False and time=False: error, 
       if units=True: 1 dim numpy.ndarray; time series,
       if time=True:  1 dim numpy.ndarray; series of single unit variances across time,
       if units=True and time=True: float; mean of single unit variances across time

    **Examples**: 
       >>> variance(np.array([[1,2,3],[4,5,6]]),units=True)
       Out[1]: np.array([ 2.25,  2.25,  2.25])
       >>> variance(np.array([[1,2,3],[4,5,6]]),time=True)
       Out[1]: np.array([ 0.66666667,  0.66666667])
       >>> variance(np.array([[1,2,3],[4,5,6]]),units=True,time=True)
       Out[1]: 0.66666666666666663

    '''

    assert(units is not False or time is not False)
    if units is True and time is False:
        return np.var(data, axis=0)
    elif units is False and time is True:
        return np.var(data, axis=1)
    elif units is True and time is True:
        return np.mean(np.var(data, axis=1))


def compound_variance(data):
    '''
    Compute the variance of the compound/sum signal.
    data is first summed across units, then the variance across time is calculated.
    
    **Args**:
       data: numpy.ndarray; 1st axis unit, 2nd axis time

    **Return**:
       float; variance across time of compound/sum signal

    **Examples**: 
       >>> compound_variance(np.array([[1,2,3],[4,5,6]]))
       Out[1]: 2.6666666666666665

    '''

    return np.var(np.sum(data, axis=0))


def spectrogram(data, tbin, twindow, Df=None, units=False, N=None, measure='power'):
    '''Calculate (smoothed) spectrogram of data. If units is True, power
    spectra are averaged across units.

    Parameters:
    -----------
    data: binned timeseries
    tbin: bin size
    twindow: size of window use for spectra
    Df: width of smoothing kernel
    units: if True, average over units
    N: population size, if not given, calculated from data
    measure: define the measure to be used (power, cross, compound_power)
    '''

    steps_window = int(np.floor(twindow / tbin))
    n_windows = int(np.floor(1. * len(data[0]) / steps_window))
    sg = []