__all__ = ['ilr',

           'rank_intensity',

           'quantile',

           'regress_poly',

           'regress_components',

           'get_average_of_timeseries',

           'compcor',

           'bandpass_filter_matrix' ]

import numpy as np

from numpy.polynomial import Legendre

from scipy import linalg

from scipy.stats import pearsonr

from scipy.stats import rankdata

import pandas as pd

from pandas import DataFrame

import statsmodels.api as sm

import statsmodels.formula.api as smf

import ants

from ants.decorators import image_method

def rank_intensity( x, mask=None, get_mask=True, method='max',  ):

    """

    Rank transform the intensity of the input image with or without masking.

    Intensities will transform from [0,1,2,55] to [0,1,2,3] so this may not be

    appropriate for quantitative images - however, you never know.  rank

    transformations generally improve robustness so it is an empirical question

    that should be evaluated.

    Arguments

    ---------

    x : ANTsImage

        input image

    mask : ANTsImage

        optional mask

    get_mask: boolean

        will estimate a mask when none provided

    method : a scipy rank method (max,min,average,dense)

    return: transformed image

    Example

    -------

    >>> img = ants.image_read(ants.get_data('r16'))

    >>> ants.rank_intensity(img)

    """

    if mask is not None:

        fir = rankdata( (x*mask).numpy(), method=method )

    elif mask is None and get_mask == True:

        mask = ants.get_mask( x )

        fir = rankdata( (x*mask).numpy(), method=method )

    else:

        fir = rankdata( x.numpy(), method=method )

    fir = fir - 1

    fir = fir.reshape( x.shape )

    rimg = ants.from_numpy( fir.astype(float)  )

    rimg = ants.iMath(rimg,"Normalize")

    ants.copy_image_info( x, rimg )

    if mask is not None:

        rimg = rimg * mask

    return( rimg )

def ilr( data_frame, voxmats, ilr_formula, verbose = False ):

    """

    Image-based linear regression.

    This function simplifies calculating p-values from linear models

    in which there is a similar formula that is applied many times

    with a change in image-based predictors.  Image-based variables

    are stored in the input matrix list. They should be named

    consistently in the input formula and in the image list.  If they

    are not, an error will be thrown.  All input matrices should have

    the same number of rows and columns.

    This function takes advantage of statsmodels R-style formulas.

    ANTsR function: `ilr`

    Arguments

    ---------

    data_frame: This data frame contains all relevant predictors except for

        the matrices associated with the image variables.  One should convert

        any categorical predictors ahead of time using `pd.get_dummies`.

    voxmats: The named list of matrices that contains the changing

        predictors.

    ilr_formula: This is a character string that defines a valid regression

        formula in the R-style.

    verbose:  will print a little bit of diagnostic information that allows

        a degree of model checking

    Returns

    -------

    A list of different matrices that contain names derived from the

    formula and the coefficients of the regression model.  The size of

    the output values ( p-values, t-values, parameter values ) will match

    the input matrix and, as such, can be converted to an image via `make_image`

    Example

    -------

    >>> nsub = 20

    >>> mu, sigma = 0, 1

    >>> outcome = np.random.normal( mu, sigma, nsub )

    >>> covar = np.random.normal( mu, sigma, nsub )

    >>> mat = np.random.normal( mu, sigma, (nsub, 500 ) )

    >>> mat2 = np.random.normal( mu, sigma, (nsub, 500 ) )

    >>> data = {'covar':covar,'outcome':outcome}

    >>> df = pd.DataFrame( data )

    >>> vlist = { "mat1": mat, "mat2": mat2 }

    >>> myform = " outcome ~ covar * mat1 "

    >>> result = ants.ilr( df, vlist, myform)

    >>> myform = " mat2 ~ covar + mat1 "

    >>> result = ants.ilr( df, vlist, myform)

    """

    nvoxmats = len( voxmats )

    if nvoxmats < 1 :

        raise ValueError('Pass at least one matrix to voxmats list')

    keylist = list(voxmats.keys())

    firstmat = keylist[0]

    voxshape = voxmats[firstmat].shape

    nvox = voxshape[1]

    nmats = len( keylist )

    for k in keylist:

        if voxmats[firstmat].shape != voxmats[k].shape:

            raise ValueError('Matrices must have same number of rows (samples)')

    # test voxel

    vox = 0

    nrows = data_frame.shape[0]

    data_frame_vox = data_frame.copy()

    for k in range( nmats ):

        data = {keylist[k]: np.random.normal(0,1,nrows) }

        temp = pd.DataFrame( data )

        data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )

    mod = smf.ols(formula=ilr_formula, data=data_frame_vox )

    res = mod.fit()

    modelNames = res.model.exog_names

    if verbose:

        print( data_frame_vox )

        print(res.summary())

    nOutcomes = len( modelNames )

    tValsOut = list()

    pValsOut = list()

    bValsOut = list()

    for k in range( len( modelNames ) ):

        bValsOut.append( np.zeros( nvox ) )

        pValsOut.append( np.zeros( nvox ) )

        tValsOut.append( np.zeros( nvox ) )

    data_frame_vox = data_frame.copy()

    for v in range( nmats ):

        data = {keylist[v]: voxmats[keylist[v]][:,k] }

        temp = pd.DataFrame( data )

        data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )

    for k in range( nvox ):

        # first get the correct data frame

        for v in range( nmats ):

            data_frame_vox[ keylist[v] ] = voxmats[keylist[v]][:,k]

        # then get the local model results

        mod = smf.ols(formula=ilr_formula, data=data_frame_vox )

        res = mod.fit()

        tvals = res.tvalues

        pvals = res.pvalues

        bvals = res.params

        for v in range( len( modelNames ) ):

            bValsOut[v][k] = bvals.iloc[v]

            pValsOut[v][k] = pvals.iloc[v]

            tValsOut[v][k] = tvals.iloc[v]

    bValsOutDict = { }

    tValsOutDict = { }

    pValsOutDict = { }

    for v in range( len( modelNames ) ):

        bValsOutDict[ 'coef_' + modelNames[v] ] = bValsOut[v]

        tValsOutDict[ 'tval_' + modelNames[v] ] = tValsOut[v]

        pValsOutDict[ 'pval_' + modelNames[v] ] = pValsOut[v]

    return {

        'modelNames': modelNames,

        'coefficientValues': bValsOutDict,

        'pValues': pValsOutDict,

        'tValues': tValsOutDict }

@image_method

def quantile(image, q, nonzero=True):

    """

    Get the quantile values from an ANTsImage

    Examples

    --------

    >>> img = ants.image_read(ants.get_data('r16'))

    >>> ants.quantile(img, 0.5)

    >>> ants.quantile(img, (0.5, 0.75))

    """

    img_arr = image.numpy()

    if isinstance(q, (list,tuple)):

        q = [qq*100. if qq <= 1. else qq for qq in q]

        if nonzero:

            img_arr = img_arr[img_arr>0]

        vals = [np.percentile(img_arr, qq) for qq in q]

        return tuple(vals)

    elif isinstance(q, (float,int)):

        if q <= 1.:

            q = q*100.

        if nonzero:

            img_arr = img_arr[img_arr>0]

        return np.percentile(img_arr[img_arr>0], q)

    else:

        raise ValueError('q argument must be list/tuple or float/int')

def regress_poly(degree, data, remove_mean=True, axis=-1):

    """

    Returns data with degree polynomial regressed out.

    :param bool remove_mean: whether or not demean data (i.e. degree 0),

    :param int axis: numpy array axes along which regression is performed

    """

    timepoints = data.shape[0]

    # Generate design matrix

    X = np.ones((timepoints, 1))  # quick way to calc degree 0

    for i in range(degree):

        polynomial_func = Legendre.basis(i + 1)

        value_array = np.linspace(-1, 1, timepoints)

        X = np.hstack((X, polynomial_func(value_array)[:, np.newaxis]))

    non_constant_regressors = X[:, :-1] if X.shape[1] > 1 else np.array([])

    betas = np.linalg.pinv(X).dot(data)

    if remove_mean:

        datahat = X.dot(betas)

    else:  # disregard the first layer of X, which is degree 0

        datahat = X[:, 1:].dot(betas[1:, ...])

    regressed_data = data - datahat

    return regressed_data, non_constant_regressors

def regress_components( data, components, remove_mean=True ):

    """

    Returns data with components regressed out.

    :param bool remove_mean: whether or not demean data (i.e. degree 0),

    :param int axis: numpy array axes along which regression is performed

    """

    timepoints = data.shape[0]

    betas = np.linalg.pinv(components).dot(data)

    if remove_mean:

        datahat = components.dot(betas)

    else:  # disregard the first layer of X, which is degree 0

        datahat = components[:, 1:].dot(betas[1:, ...])

    regressed_data = data - datahat

    return regressed_data

def get_average_of_timeseries( image, idx=None ):

    """Average the timeseries into a dimension-1 image.

    image: input time series image

    idx:  indices over which to average

    """

    imagedim = image.dimension

    if idx is None:

        idx = range( image.shape[ imagedim - 1 ] )

    i0 = ants.slice_image( image, axis=image.dimension-1, idx=idx[0] ) * 0

    wt = 1.0 / len( idx )

    for k in idx:

        i0 = i0 + ants.slice_image( image, axis=image.dimension-1, idx=k ) * wt

    return( i0 )

def bandpass_filter_matrix( matrix,

    tr=1, lowf=0.01, highf=0.1, order = 3):

    """

    Bandpass filter the input time series image

    ANTsR function: `frequencyFilterfMRI`

    Arguments

    ---------

    image: input time series image

    tr:    sampling time interval (inverse of sampling rate)

    lowf:  low frequency cutoff

    highf: high frequency cutoff

    order: order of the butterworth filter run using `filtfilt`

    Returns

    -------

    filtered matrix

    Example

    -------

    >>> import numpy as np

    >>> import ants

    >>> import matplotlib.pyplot as plt

    >>> brainSignal = np.random.randn( 400, 1000 )

    >>> tr = 1

    >>> filtered = ants.bandpass_filter_matrix( brainSignal, tr = tr )

    >>> nsamples = brainSignal.shape[0]

    >>> t = np.linspace(0, tr*nsamples, nsamples, endpoint=False)

    >>> k = 20

    >>> plt.plot(t, brainSignal[:,k], label='Noisy signal')

    >>> plt.plot(t, filtered[:,k], label='Filtered signal')

    >>> plt.xlabel('time (seconds)')

    >>> plt.grid(True)

    >>> plt.axis('tight')

    >>> plt.legend(loc='upper left')

    >>> plt.show()

    """

    from scipy.signal import butter, filtfilt

    def butter_bandpass(lowcut, highcut, fs, order ):

        nyq = 0.5 * fs

        low = lowcut / nyq

        high = highcut / nyq

        b, a = butter(order, [low, high], btype='band')

        return b, a

    def butter_bandpass_filter(data, lowcut, highcut, fs, order ):

        b, a = butter_bandpass(lowcut, highcut, fs, order=order)

        y = filtfilt(b, a, data)

        return y

    fs = 1/tr   # sampling rate based on tr

    nsamples = matrix.shape[0]

    ncolumns = matrix.shape[1]

    matrixOut = matrix.copy()

    for k in range( ncolumns ):

        matrixOut[:,k] = butter_bandpass_filter(

            matrix[:,k], lowf, highf, fs, order=order )

    return matrixOut

def clean_data(arr, standardize=True):

    """

    Remove columns from a NumPy array that have no variation or contain NA/Inf values.

    Optionally standardize the remaining data.

    :param arr: NumPy array to be cleaned.

    :param standardize: Boolean, if True standardize the data.

    :return: Cleaned (and optionally standardized) NumPy array.

    """

    valid_columns = []

    for i in range(arr.shape[1]):

        column = arr[:, i]

        if np.any(column != column[0]) and not np.any(np.isnan(column)) and not np.any(np.isinf(column)):

            valid_columns.append(i)

    cleaned_data = arr[:, valid_columns]

    if standardize:

        mean = np.mean(cleaned_data, axis=0)

        std_dev = np.std(cleaned_data, axis=0)

        # Avoid division by zero in case of zero standard deviation

        std_dev[std_dev == 0] = 1

        cleaned_data = (cleaned_data - mean) / std_dev

    return cleaned_data

def compcor( boldImage, ncompcor=4, quantile=0.975, mask=None, filter_type=False, degree=2 ):

    """

    Compute noise components from the input image

    ANTsR function: `compcor`

    this is adapted from nipy code https://github.com/nipy/nipype/blob/e29ac95fc0fc00fedbcaa0adaf29d5878408ca7c/nipype/algorithms/confounds.py

    Arguments

    ---------

    boldImage: input time series image

    ncompcor:  number of noise components to return

    quantile:  quantile defining high-variance

    mask:      mask defining brain or specific tissues

    filter_type: type off filter to apply to time series before computing

                 noise components.

        'polynomial' - Legendre polynomial basis

        False - None (mean-removal only)

    degree: order of polynomial used to remove trends from the timeseries

    Returns

    -------

    dictionary containing:

        components: a numpy array

        basis: a numpy array containing the (non-constant) filter regressors

    Example

    -------

    >>> cc = ants.compcor( ants.image_read(ants.get_ants_data("ch2")) )

    """

    def compute_tSTD(M, quantile, x=0, axis=0):

        stdM = np.std(M, axis=axis)

        # set bad values to x

        stdM[stdM == 0] = x

        stdM[np.isnan(stdM)] = x

        tt = round( quantile*100 )

        threshold_std = np.percentile( stdM, tt )

    #    threshold_std = quantile( stdM, quantile )

        return { 'tSTD': stdM, 'threshold_std': threshold_std}

    if mask is None:

        temp = ants.slice_image( boldImage, axis=boldImage.dimension-1, idx=0 )

        mask = ants.get_mask( temp )

    imagematrix = ants.timeseries_to_matrix( boldImage, mask )

    temp = compute_tSTD( imagematrix, quantile, 0 )

    tsnrmask = ants.make_image( mask,  temp['tSTD'] )

    tsnrmask = ants.threshold_image( tsnrmask, temp['threshold_std'], temp['tSTD'].max() )

    M = ants.timeseries_to_matrix( boldImage, tsnrmask )

    components = None

    basis = np.array([])

    if filter_type in ('polynomial', False):

        M, basis = regress_poly(degree, M)

#    M = M / compute_tSTD(M, 1.)['tSTD']

    # "The covariance matrix C = MMT was constructed and decomposed into its

    # principal components using a singular value decomposition."

    M = clean_data( M, standardize=True )

    u, _, _ = linalg.svd(M, full_matrices=False)

    if components is None:

        components = u[:, :ncompcor]

    else:

        components = np.hstack((components, u[:, :ncompcor]))

    if components is None and ncompcor > 0:

        raise ValueError('No components found')

    return { 'components': components, 'basis': basis }