# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_utils.ipynb.

# %% auto 0
__all__ = ['make_stable', 'add_stability', 'add_non_const', 'apply_log2_fold', 'apply_signed_correlation',
           'apply_standard_scaler', 'kpss_test']

# %% ../nbs/01_utils.ipynb 3
import pandas as pd, numpy as np
import seaborn as sns

from typing import List, Optional, Union

from sklearn.preprocessing import StandardScaler
from statsmodels.tsa.stattools import kpss

# %% ../nbs/01_utils.ipynb 5
from atyp import SeriesLike, DataFrame, NPArray

# %% ../nbs/01_utils.ipynb 6
from gnr.static import (
    LOG_SAFETY, SSR_CHI2TEST, LAG_ORDER,
    DEFAULT_TEST, DEFAULT_GRANGER_CAUSALITY_TESTS,
    MOCK_GENES
)

# %% ../nbs/01_utils.ipynb 8
def _add_suffix(arr, suffix:str):
    return list(map(lambda e: f'{e}_{suffix}', arr))

def make_stable(
    df: DataFrame, 
    copy: Optional[bool] = True,
    add_perm: Optional[bool] = True,
    use_rand: Optional[bool] = False,
    no_const: Optional[bool] = True
) -> DataFrame:    
    df_safe = (df.copy() if copy else df)

    if add_perm:
        df_safe = df_safe + LOG_SAFETY
    
    # NOTE: to ensure that no column in df is constant
    if use_rand or no_const:
        rng_mat = np.random.randn(*df.shape) * LOG_SAFETY
        df_safe = df_safe + rng_mat

    return df_safe
    
def add_stability(df):
    return make_stable(df, copy=False, add_perm=True)

def add_non_const(df):
    return make_stable(df, copy=False, add_perm=False, use_rand=True)
    
def apply_log2_fold(df_granger: DataFrame):
    return -np.log2(df_granger + (2 ** -20))    

def apply_signed_correlation(df_tseries: DataFrame, df_granger: DataFrame) -> DataFrame:
    var_names = df_tseries.index.values
    signed = -np.sign(df_tseries.T.corr()) * np.log(add_stability(df_granger.values))
    df_res = pd.DataFrame(signed, index=var_names, columns=var_names)
    return df_res

def apply_standard_scaler(df: DataFrame, return_scaler: Optional[bool] = False) -> DataFrame:
    scaler = StandardScaler()
    scaler = scaler.fit(df.T)
    df_res = pd.DataFrame(scaler.transform(df.T)).T
    df_res.index = df.index
    if return_scaler:
        df_res, scaler
    return df_res

# %% ../nbs/01_utils.ipynb 10
def kpss_test(series, print: Optional[bool] = False, **kwargs) -> None:
    # KPSS Null hypothesis: there is a no unit root, meaning series is stationary
    statistic, p_value, n_lags, critical_values = kpss(series, **kwargs)

    # Format Output    
    result_str = ''
    result_str += f'KPSS Statistic: {statistic}\n'
    result_str += f'p-value: {p_value}\n'
    result_str += f'num lags: {n_lags}\n'
    result_str += f'Critial Values:\n'    
    for key, value in critical_values.items():
        result_str += f'\t{key}:{value}\n'
    result_str += f'Result: The series is {"not " if p_value < 0.05 else ""}stationary'

    if print:
        print(result_str)
    return

# %% ../nbs/01_utils.ipynb 12
def _prep_args_for_granger_causality_tests(
    test: Optional[str] = DEFAULT_TEST, 
    lag_order: Optional[int] = LAG_ORDER,
    max_lag: Optional[tuple] = (LAG_ORDER, ),    
):
    '''
    Parameters
    ----------
    test : str, default='ssr_chi2test'
        the kind of statistical test to use
    
    lag_order : int, default=1
        how long to lag
    
    max_lag : tuple, optional
        if `None` coerced to `(1, )`
    
    Returns
    -------
    test : str
        `'ssr_chi2test'` if `test` was `None`
        
    lag_order : int
        `1` if `lag_order` was `None`

    max_lag : tuple
        `(1, )` if `max_lag` was `None`

    Notes
    -----
    lag_order
        `lag_order = 1` # becuase we got this value before. We are not suppose to add 1 to it
    
    max_lag
        `max_lag = (lag_order, )`
    '''
    if test is None:
        test = SSR_CHI2TEST

    if lag_order is None:
        lag_order = LAG_ORDER

    if max_lag is None:
        max_lag = (lag_order,)
    
    return test, lag_order, max_lag

def _prep_vars_for_granger_causality_tests(
    df: DataFrame, 
    x_vars: Optional[SeriesLike] = None, 
    y_vars: Optional[SeriesLike] = None,
    do_safety_check: Optional[bool] = True     
):
    '''
    Ensures `df` is safe for Granger Causality tests and that 
    `x_vars` and `y_vars` are not `None`.

    Parameters
    ----------
    df : pd.DataFrame
        Pandas DataFrame where rows are the response variable (genes), and 
        columns are predictors (expression).
    
    x_vars : SeriesLike, optional
        A subset of response variable (genes) to compute granger's causality test with. 
        If not provided, defaults to `df.index.values` i.e. all rows in `df`.


    y_vars : SeriesLike, optional
        A subset of response variable (genes) to compute granger's causality test with. 
        If not provided, defaults to `df.index.values` i.e. all rows in `df`.

    do_safety_check : bool, default=True
        Whether or not to check if there is a constant column in `df`.

        
    Returns
    -------
    df : pd.DataFrame
        input `df` possibly with small random noise added to prevent constant columns

    x_vars : SeriesLike
        input `x_vars` or `df.index.values` if `x_vars` was `None`

    y_vars : SeriesLike
        input `y_vars` or `df.index.values` if `y_vars` was `None`

    Notes
    -----
    If `x_var = ['A', B']` and `y_var = ['C', 'D]', then when calling Granger's 
    Causality the matrix returned has shape:
    ```
        | C | D |
        ---------
    | A |   |   |
    -------------
    | B |   |   |

    ```
    '''
    if do_safety_check:
        df = add_non_const(df)

    if x_vars is None:
        x_vars = df.index.values

    if y_vars is None:
        y_vars = df.index.values
    
    return df, x_vars, y_vars
