import numpy as np
from scipy.stats import t

def mvtstat(data):
    """ A function to compute the multivariate t-statistic

    Parameters
    -----------------
    data:  numpy.ndarray of shape (Dim, nsubj)
          Here Dim is the size of the field and nsubj is the number of subjects

    Returns
    -----------------
    tstat:   numpy.ndarray of shape (Dim)
          Each entry is the is the t-statistic calulcated across subjects
    mean:    numpy.ndarray of shape (Dim)
          Each entry is the is the mean calulcated across subjects
    std:     numpy.ndarray of shape (Dim)
          Each entry is the is the standard deviation calulcated across subjects

    Examples
    -----------------
    # tstat of random noise
    noise = np.random.randn(50,50,20); arrays = mvtstat(noise);  tstat = arrays[0]
    # For comparison to MATLAB
    a = np.arange(12).reshape((3,4)).transpose()+1; tstat = mvtstat(a)[0]
    """
    # Obtain the size of the array
    s_data = np.shape(data)

    # Obtain the dimensions
    dim = s_data[0:-1]

    # Obtain the number of dimensions
    n_dim = len(dim)

    # Obtain the number of subjects
    nsubj = s_data[-1]

    # Get the mean and stanard deviation along the number of subjects
    # Remember in python D is the last dimension of a D+1 array
    xbar = data.mean(n_dim)

    # Calculate the standard deviation (multiplying to ensure the population std is used!)
    std_dev = data.std(n_dim)*np.sqrt((nsubj/(nsubj-1.)))

    # Compute Cohen's d
    cohensd = xbar/std_dev
    tstat = np.sqrt(nsubj)*cohensd

    return tstat, xbar, std_dev

def tstat2pval(tstat, degrees_of_freedom, do2sample=1):
    """
    tstat2pval(tstat, degrees_of_freedom, do2sample) calculates the p-values from the
    t-statistic.
    
    Parameters
    ----------
    tstat : array_like
        An array of t-statistics
    df : int or float
        Degrees of freedom
    do2sample : int, optional
        0/1, 1: compute two sample p-values. 0: compute one sample
        p-values. Default is 1 (two sample).
    
    Returns
    -------
    pvals : ndarray
        Array of p-values corresponding to the input t-statistics
    
    Examples
    --------
    >>> # Two-sample test (default)
    >>> tstat = [2.5, -1.8, 3.2]
    >>> df = 10
    >>> pvals = tstat2pval(tstat, df)
    
    >>> # One-sample test  
    >>> pvals = tstat2pval(tstat, df, do2sample=0)
    
    Author: Samuel Davenport (converted to Python)
    """
    
    # Convert input to numpy array for consistent handling
    tstat = np.asarray(tstat)
    
    # Main calculation
    if do2sample == 1:
        # Two-sample test: two-tailed p-value
        pvals = 2 * (1 - t.cdf(np.abs(tstat), degrees_of_freedom))
    else:
        # One-sample test: one-tailed p-value
        pvals = 1 - t.cdf(tstat, degrees_of_freedom)
    
    return pvals