def flatten(foo):
    # Taken from https://stackoverflow.com/a/5286571
    for x in foo:
        if hasattr(x, '__iter__') and not isinstance(x, str):
            for y in flatten(x):
                yield y
        else:
            yield x

def freedman_lane(data_df, Yvar, Xvars, Zvars, n_perms=10000, stat='tstat', perm_func=None, perm_func_args=None, 
           surrogates=None, return_null=False, return_surrogates=False):
    """
    Use permutation testing (via random shuffling or a user provided function) to estimate the significance of an EV (Xvar)
    in a multiple linear regression (Yvar ~ Xvar + Zvars) while "correcting for" (regressing-out the estimated effects of)
    other covariates (Zvars).

    The null distribution of test statistics is generated by permuting the residuals of a reduced model (Yvar ~ Zvars)
    (following Freedman and Lane, 1983; DOI: 10.2307/1391660).

    Note:
    Interactions are not curently supported, and categorical variables are currently only supported when stat='fstat'.

    Parameters:
    -----------
    data_df : Pandas DataFrame
        A data frame containing data to be modeled.
    Yvar : str
        Name of the DataFrame column containing the dependent (endog) variable
    Xvar : str
        Name of the DataFrame column containing the independent (exog) variable of interest.
    Zvars : str or list or str
        Name of the DataFrame column(s) containing the independent (exog) variable(s) to use as a covariate.
    n_perms : int, optional
        Number of surrogate data sets to generate. Default is 10,000.
    stat : string, optional
        Which test statistic to evaluate significance for and return. 'tstat' evaluates the significant of the t-statistic
        for the effect of Xvar. 'pcorr' evaluates significance of the partial pearson correlation coefficient between the 
        residuals of Yvar ~ Zvars and Xvar ~ Zvars. 'tstat' and 'pcorr' will always return the same p-value (because the
        partial correlation is calculated from the t-statistic); I provide both for convenience. 'fstat' uses a Type II ANOVA 
        to evaluate the significance of the main effect for Xvar.
    perm_func : function, optional
        Custom function to generate surrogate data. Must accept a 1D array (the data to be permuted) as the first argument,
        and n_perms (an integer) as the second argument. If no custom function is provided, surrogate data will be generated
        via random shuffling (which assumes full exchangability).
    perm_func_args : dict, optional
        Dictionary containing additional arguments (including potentially additional data) to pass to perm_func.
    surrogates: ndarray, optional
        Surrogate data sets generated and returned by a previous run of lmperm which evaluated a model with the same
        Yvar and Zvars as currently specified. 
    return_null : bool
        Return the null distribution of statistic values. Useful if you want to plot the observed value against the null
        distribution. Defaults to False.
    return_surrogates : bool
        Return an array containing surrogate data sets. Useful time saver if you want to test the relationship of a 
        given Yvar to multiple Xvars while correcting for the same Zvars. Defaults to False.

    Returns:
    --------
    stat_observed : float
        The test statistic for Xvar from the full, un-permuted model.
    pvals : Pandas Series
        The probability of the observed test statistic relative to the null distribution.
        'p_greater' is Pr(stat_observed >= null). 'p_less' is Pr(stat_observed <= null).
        'p_abs' is Pr(|stat_observed| >= |null|). If stat='fstat', only p_greater is returned.
    fit_Full : statsmodels RegressionResults
        Results (coefficients, fit measures, p-values, etc.) of the parametric OLS fit to to the full model. 
    stats_null : ndarray, optional
        Null distribution of test statistics generated via permutation testing.
    resid_Reduced_Perms : ndarray, optional
        Permuted residuals of the reduced model (Yvar ~ Zvars).

    
    """
    import warnings
    import numpy as np
    import pandas as pd
    import statsmodels.api as sm
    import statsmodels.formula.api as smf
    from statsmodels.stats import anova

    # Check inputs.
    # This is not Pythonic, but is more user friendly than letting exceptions raise down stream.
    # Some interesting ideas here on how to make checking more Pythonic:
    # https://stackoverflow.com/questions/19684434/best-way-to-check-function-arguments
    assert (isinstance(data_df, pd.core.frame.DataFrame)), "data_df must be a Pandas DataFrame!"
    assert (type(Yvar) == str), "Yvar must be a string!"
    assert (type(Xvars) == str or type(Xvars) == list), "Xvars must be a string or a list!"
    assert (type(Zvars) == str or type(Zvars) == list), "Zvars must be a string or a list!"
    assert (type(n_perms) == int), "n_perms must be an integer!"
    if perm_func is not None:
        assert (callable(perm_func)), "perm_func must be a function!"
    if perm_func_args is not None:
        assert (type(perm_func_args) == dict), "perm_func_args must be a dictionary!"
    if surrogates is not None:
        assert (type(surrogates) == np.ndarray), "surrogates must be an ndarray!"

    # Generate a formula string for Zvars.
    if type(Zvars) == list:
        Zvars_formula_string = " + ".join(Zvars)
        Zsize = len(Zvars)
    elif type(Zvars) == str:
        Zvars_formula_string = Zvars
        Zsize = 1

    # Generate a formula string for Xvars
    if type(Xvars) == list:
        Xvars_formula_string = " + ".join(Xvars)
        if stat == 'fstat' is False:
            warnings.warn("Multiple Xvars provided; setting stat='fstat'.")
            stat = 'fstat'
    elif type(Xvars) == str:
        Xvars_formula_string = Xvars

    # Fit the full model with observed data.
    formula_Full = f"{Yvar} ~ {Xvars_formula_string} + {Zvars_formula_string}"
    fit_Full = smf.ols(formula_Full, data=data_df).fit()
    if stat == 'fstat':
        contrast_Xvars_string = " = ".join(Xvars)
        contrast_formula = f"{contrast_Xvars_string} = 0"
        f_res = fit_Full.f_test(contrast_formula)
        stat_observed = f_res.fvalue[0][0]
    elif stat == 'pcorr':
        stat_observed = (fit_Full.tvalues[Xvars] / np.sqrt(fit_Full.tvalues[Xvars]**2 + fit_Full.df_resid))
    elif stat == 'tstat':
        stat_observed = fit_Full.tvalues[Xvars]    

    # Fit the reduced model with the observed data.
    formula_Reduced = f"{Yvar} ~ {Zvars_formula_string}"
    model_Reduced = smf.ols(formula_Reduced, data=data_df)
    fit_Reduced = model_Reduced.fit()
    resid_Reduced = fit_Reduced.resid.values
    Yhat_reduced = model_Reduced.predict(fit_Reduced.params)
        
    # If the user has provided previously generated surrogates, use those.
    if surrogates is not None:
        resid_Reduced_perms = surrogates
    # Otherwise, generate surrogate data sets.
    # TO-DO: Generate permutations in parallel via multiprocessing/joblib/ray.
    else:
        # Permute the residuals from the reduced model.
        # If a custom permutation function is provided, use that.
        if perm_func:
            resid_Reduced_perms = perm_func(resid_Reduced, n_perms, **perm_func_args)
        # Otherwise, just do random shuffling.
        else:
            resid_Reduced_perms = []
            for i in range(n_perms):
                resid_Reduced_perms.append(np.random.permutation(resid_Reduced))
            resid_Reduced_perms = np.array(resid_Reduced_perms)

    # Generate a null distribution by calculating test statistics from the full model for each of the permutations.
    stats_null = []
    exog_cols = []
    exog_cols.append(Xvars)
    exog_cols.append(Zvars)
    exog_cols = list(flatten(exog_cols))
    ev_array = data_df[exog_cols].values
    # The non-formula OLS interface is much faster but does not automatically add an intercept, so we must add it ourselves.
    intercept = np.reshape(np.ones_like(data_df[exog_cols[0]]),(ev_array.shape[0],1))
    ev_array = np.concatenate((intercept, ev_array), axis=1)
    for perm_iter in resid_Reduced_perms:
        Ypi = perm_iter + Yhat_reduced
        fit_Perm = sm.OLS(Ypi, ev_array).fit()
        if stat == 'fstat':
            # Create a contrast array that includes all regressors
            contrast_array = np.identity(len(fit_Perm.params))
            # Drop the intercept
            contrast_array = contrast_array[1:,:]
            # Drop Zvars
            contrast_array = contrast_array[:-Zsize,:]
            f_res = fit_Perm.f_test(contrast_array)
            stats_null.append(f_res.fvalue[0][0])
        else:
            if stat == 'tstat':
                stats_null.append(fit_Perm.tvalues[1])
            elif stat == 'pcorr':
                stats_null.append(fit_Perm.tvalues[1] / np.sqrt(fit_Perm.tvalues[1]**2 + fit_Perm.df_resid))
    stats_null = np.array(stats_null)

    # Calculate the probability of the observed test statistic relative to the null distribution.
    if stat == 'fstat':
        p_greater = (np.sum(stats_null >= stat_observed) + 1) / (n_perms + 1)
        pvals = pd.Series([p_greater], index=['p_greater'])
    elif stat == 'pcorr' or stat == 'tstat':
        p_greater = (np.sum(stats_null >= stat_observed) + 1) / (n_perms + 1)
        p_less = (np.sum(stats_null <= stat_observed) + 1) / (n_perms + 1)
        p_abs = (np.sum(np.abs(stats_null) >= np.abs(stat_observed)) + 1) / (n_perms + 1)
        pvals = pd.Series([p_greater, p_less, p_abs], index=['p_greater', 'p_less', 'p_abs'])

    # Collect outputs.
    out = (stat_observed, pvals, fit_Full)
    if return_null:
        out += (stats_null,)
    if return_surrogates:
        out += (resid_Reduced_perms,)

    return out

def manly(data_df, Yvar, Xvars, Zvars, stat='tstat', n_perms=10000, perm_func=None, perm_func_args=None, 
           surrogates=None, return_null=False, return_surrogates=False):
    """
    Use permutation testing (via random shuffling or a user provided function) to estimate the significance of an EV (Xvar)
    in a multiple linear regression (Yvar ~ Xvar + Zvars) while "correcting for" (regressing-out the estimated effects of)
    other covariates (Zvars).

    The null distribution of test statistics is generated by permuting Yvar (following Manly, 1986; DOI: 10.1007/BF02515450).

    Note:
    Interactions are not curently supported, and categorical variables are currently only supported when stat='fstat'.

    Parameters:
    -----------
    data_df : Pandas DataFrame
        A data frame containing data to be modeled.
    Yvar : str
        Name of the DataFrame column containing the dependent (endog) variable
    Xvar : str
        Name of the DataFrame column containing the independent (exog) variable of interest.
    Zvars : str or list or str
        Name of the DataFrame column(s) containing the independent (exog) variable(s) to use as a covariate.
    n_perms : int, optional
        Number of surrogate data sets to generate. Default is 10,000.
    stat : string, optional
        Which test statistic to evaluate significance for and return. 'tstat' evaluates the significant of the t-statistic
        for the effect of Xvar. 'pcorr' evaluates significance of the partial pearson correlation coefficient between the 
        residuals of Yvar ~ Zvars and Xvar ~ Zvars. 'tstat' and 'pcorr' will always return the same p-value (because the
        partial correlation is calculated from the t-statistic); I provide both for convenience. 'fstat' uses a Type II ANOVA 
        to evaluate the significance of the main effect for Xvar. 
    perm_func : function, optional
        Custom function to generate surrogate data. Must accept a 1D array (the data to be permuted) as the first argument,
        and n_perms (an integer) as the second argument. If no custom function is provided, surrogate data will be generated
        via random shuffling (which assumes full exchangability).
    perm_func_args : dict, optional
        Dictionary containing additional arguments (including potentially additional data) to pass to perm_func.
    surrogates: ndarray, optional
        Surrogate data sets generated and returned by a previous run of lmperm which evaluated a model with the same
        Yvar and Zvars as currently specified. 
    return_null : bool
        Return the null distribution of statistic values. Useful if you want to plot the observed value against the null
        distribution. Defaults to False.
    return_surrogates : bool
        Return an array containing surrogate data sets. Useful time saver if you want to test the relationship of a 
        given Yvar to multiple Xvars while correcting for the same Zvars. Defaults to False.

    Returns:
    --------
    stat_observed : float
        The test statistic for Xvar from the full, un-permuted model.
    pvals : Pandas Series
        The probability of the observed test statistic relative to the null distribution.
        'p_greater' is Pr(stat_observed >= null). 'p_less' is Pr(stat_observed <= null).
        'p_abs' is Pr(|stat_observed| >= |null|). If stat='fstat', only p_greater is returned.
    fit_Full : statsmodels RegressionResults
        Results (coefficients, fit measures, p-values, etc.) of the parametric OLS fit to to the full model. 
    stats_null : ndarray, optional
        Null distribution of test statistics generated via permutation testing.
    resid_Reduced_Perms : ndarray, optional
        Permuted residuals of the reduced model (Yvar ~ Zvars).

    
    """
    import warnings
    import numpy as np
    import pandas as pd
    import statsmodels.api as sm
    import statsmodels.formula.api as smf
    from statsmodels.stats import anova

    # Check inputs.
    # This is not Pythonic, but is more user friendly than letting exceptions raise down stream.
    # Some interesting ideas here on how to make checking more Pythonic:
    # https://stackoverflow.com/questions/19684434/best-way-to-check-function-arguments
    assert (isinstance(data_df, pd.core.frame.DataFrame)), "data_df must be a Pandas DataFrame!"
    assert (type(Yvar) == str), "Yvar must be a string!"
    assert (type(Xvars) == str or type(Xvars) == list), "Xvars must be a string or a list!"
    assert (type(Zvars) == str or type(Zvars) == list), "Zvars must be a string or a list!"
    assert (type(n_perms) == int), "n_perms must be an integer!"
    if perm_func is not None:
        assert (callable(perm_func)), "perm_func must be a function!"
    if perm_func_args is not None:
        assert (type(perm_func_args) == dict), "perm_func_args must be a dictionary!"
    if surrogates is not None:
        assert (type(surrogates) == np.ndarray), "surrogates must be an ndarray!"

    # Generate a formula string for Zvars.
    if type(Zvars) == list:
        Zvars_formula_string = " + ".join(Zvars)
        Zsize = len(Zvars)
    elif type(Zvars) == str:
        Zvars_formula_string = Zvars
        Zsize = 1

    # Generate a formula string for Xvars
    if type(Xvars) == list:
        Xvars_formula_string = " + ".join(Xvars)
        if stat == 'fstat' is False:
            warnings.warn("Multiple Xvars provided; setting stat='fstat'.")
            stat = 'fstat'
    elif type(Xvars) == str:
        Xvars_formula_string = Xvars

    # Fit the full model with observed data.
    formula_Full = f"{Yvar} ~ {Xvars_formula_string} + {Zvars_formula_string}"
    fit_Full = smf.ols(formula_Full, data=data_df).fit()
    if stat == 'fstat':
        contrast_Xvars_string = " = ".join(Xvars)
        contrast_formula = f"{contrast_Xvars_string} = 0"
        f_res = fit_Full.f_test(contrast_formula)
        stat_observed = f_res.fvalue[0][0]
    elif stat == 'pcorr':
        stat_observed = (fit_Full.tvalues[Xvars] / np.sqrt(fit_Full.tvalues[Xvars]**2 + fit_Full.df_resid))
    elif stat == 'tstat':
        stat_observed = fit_Full.tvalues[Xvars]
    stat_observed = stat_observed.values
        
    # If the user has provided previously generated surrogates, use those.
    if surrogates is not None:
        Y_perms = surrogates
    # Otherwise, generate surrogate data sets.
    # TO-DO: Generate permutations in parallel via multiprocessing/joblib/ray.
    else:
        # Permute the Y variable.
        # If a custom permutation function is provided, use that.
        if perm_func:
            Y_perms = perm_func(data_df[Yvar].values, n_perms, **perm_func_args)
        # Otherwise, just do random shuffling.
        else:
            Y_perms = []
            for i in range(n_perms):
                Y_perms.append(np.random.permutation(data_df[Yvar].values))
            Y_perms = np.array(Y_perms)

    # Generate a null distribution by calculating test statistics from the full model for each of the permutations of Y.
    stats_null = []
    exog_cols = []
    exog_cols.append(Xvars)
    exog_cols.append(Zvars)
    exog_cols = list(flatten(exog_cols))
    ev_array = data_df[exog_cols].values
    # The non-formula OLS interface is much faster but does not automatically add an intercept, so we must add it ourselves.
    intercept = np.reshape(np.ones_like(data_df[exog_cols[0]]),(ev_array.shape[0],1))
    ev_array = np.concatenate((intercept, ev_array), axis=1)
    for perm_iter in Y_perms:
        Ypi = perm_iter
        fit_Perm = sm.OLS(Ypi, ev_array).fit()
        if stat == 'fstat':
            # Create a contrast array that includes all regressors
            contrast_array = np.identity(len(fit_Perm.params))
            # Drop the intercept
            contrast_array = contrast_array[1:,:]
            # Drop Zvars
            contrast_array = contrast_array[:-Zsize,:]
            f_res = fit_Perm.f_test(contrast_array)
            stats_null.append(f_res.fvalue[0][0])
        else:
            if stat == 'tstat':
                stats_null.append(fit_Perm.tvalues[1])
            elif stat == 'pcorr':
                stats_null.append(fit_Perm.tvalues[1] / np.sqrt(fit_Perm.tvalues[1]**2 + fit_Perm.df_resid))
    stats_null = np.array(stats_null)

    # Calculate the probability of the observed test statistic relative to the null distribution.
    if stat == 'fstat':
        p_greater = (np.sum(stats_null >= stat_observed) + 1) / (n_perms + 1)
        pvals = pd.Series([p_greater], index=['p_greater'])
    elif stat == 'pcorr' or stat == 'tstat':
        p_greater = (np.sum(stats_null >= stat_observed) + 1) / (n_perms + 1)
        p_less = (np.sum(stats_null <= stat_observed) + 1) / (n_perms + 1)
        p_abs = (np.sum(np.abs(stats_null) >= np.abs(stat_observed)) + 1) / (n_perms + 1)
        pvals = pd.Series([p_greater, p_less, p_abs], index=['p_greater', 'p_less', 'p_abs'])

    # Collect outputs.
    out = (stat_observed, pvals, fit_Full)
    if return_null:
        out += (stats_null,)
    if return_surrogates:
        out += (Y_perms,)

    return out