# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_op.ipynb.

# %% auto 0
__all__ = ['GrangerCausality']

# %% ../nbs/03_op.ipynb 3
from sklearn.base import BaseEstimator
import numpy as np, pandas as pd

import seaborn as sns

from dataclasses import dataclass, field
from typing import Optional, ClassVar

# %% ../nbs/03_op.ipynb 4
from atyp import SeriesLike, DataFrame

# %% ../nbs/03_op.ipynb 5
from gnr.static import (
    LAG_ORDER, DEFAULT_TEST,

    LOG2_FOLD, SIGNED_CORRELATION, STANDARD_SCALER, KNOWN_TRANSFORMS
)

from gnr.utils import (    
    apply_log2_fold, apply_signed_correlation, apply_standard_scaler,
    _prep_args_for_granger_causality_tests,
    _prep_vars_for_granger_causality_tests
)

from gnr.fn import (
    shift_trajectories, get_pval_from_granger_causality_tests, 
    grangers_causation_matrix, calculate_granger_causation
)

# %% ../nbs/03_op.ipynb 6
class GrangerCausality(BaseEstimator):
    '''
    Computes Granger Causality
    Check Granger Causality of all possible combinations of the Time series.

    The rows are the response variable, columns are predictors. The values in the table 
    are the P-Values. P-Values lesser than the significance level (0.05), implies 
    the Null Hypothesis that the coefficients of the corresponding past values is 
    zero, that is, the X does not cause Y can be rejected.

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

    shift : int, default=10
        number to shift `df`'s values by

    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, )`

    n_jobs : int, default=-1
        number of cpu threads to use during calculation
        
    Returns
    -------
    df_res : pd.DataFrame    
        Pandas DataFrame with shape `(len(x_vars), len(y_vars))` containing the
        minimum p-value from Granger's Causation Tests

    Notes
    -----
    If `x_var = ['A', B']` and `y_var = ['C', 'D]', the Granger's Causality matrix we
    return has shape:
    ```
        | C_y | D_y |
        -------------
    | A_x |      |     |
    --------------------
    | B_x |      |     |

    ```

    See Also
    --------
    GrangerCausality._prep_args_for_granger_causality_tests :
        validates input arguments

    GrangerCausality._prep_vars_for_granger_causality_tests :
        validates input variables
    '''
    _LOG2: ClassVar[str] = LOG2_FOLD
    _SCOR: ClassVar[str] = SIGNED_CORRELATION
    _STDN: ClassVar[str] = STANDARD_SCALER
    KNOWN_TRANSFORMS: ClassVar[str] = KNOWN_TRANSFORMS


    def __init__(
        self,
        x_vars: Optional[SeriesLike] = None, 
        y_vars: Optional[SeriesLike] = None,
        shift: Optional[int] = 10,
        test: Optional[str] = DEFAULT_TEST,
        lag_order: Optional[int] = LAG_ORDER,
        max_lag: Optional[tuple] = (LAG_ORDER, ),
        n_jobs: Optional[int] = -1
    ):        
        
    

        test, lag_order, max_lag = _prep_args_for_granger_causality_tests(test, lag_order, max_lag)

        self.x_vars = x_vars
        self.y_vars = y_vars
        self.shift = shift
        self.test = test
        self.max_lag = max_lag
        self.lag_order = lag_order
        self.n_jobs = n_jobs

        # original DataFrame
        self.df_org = None

        # results DataFrame
        self.df_res = None
        
        # scalar if transforming
        self.scaler = None

    
    def process_fit_params(self, **fit_params) -> dict:
        '''
        fit_params : dict, optional

        NOTE: you can use either of the following. The first is more explicit, the second is more concise.
        ```
            {
                "apply" : str, default=None | "log2_fold" | "signed_correlation" | "standard_scaler"
            }

            {
                "use_cached" : bool, default=False
                "log2_fold" : bool, default=False
                "signed_correlation" : bool, default=False
                "standard_scaler" : bool, default=False
            }
        ```
        if `apply` is not `None`, then the following apply:
        - `log2_fold` : will apply log2 fold change to `df_res`
        - `signed_correlation` : will apply signed correlation to `df_res`
        - `standard_scaler` : will apply standard scaler to `df_org` and `df_res`
        '''
        apply = fit_params.get('apply', None)
        if apply is not None and apply not in self.KNOWN_TRANSFORMS:
            if isinstance(apply, str):
                for app in apply.split(' '):
                    if app not in self.KNOWN_TRANSFORMS:
                        continue
                    fit_params[app] = True
        return fit_params



    def fit(self, X):
        '''
        Update internal parameters prior to computation in `transform`
        '''
        test, lag_order, max_lag = _prep_args_for_granger_causality_tests(self.test, self.lag_order, self.max_lag)
        X, x_vars, y_vars = _prep_vars_for_granger_causality_tests(X, self.x_vars, self.y_vars, do_safety_check=True)
        self.test = test
        self.lag_order = lag_order
        self.max_lag = max_lag
        self.x_vars = x_vars
        self.y_vars = y_vars
        return self
    
    def transform(self, X: DataFrame, y = None, **fit_params):   
        # Save input DataFrame
        self.df_org = X.copy()
        fit_params = self.process_fit_params(**fit_params)

        # Apply standard scaler
        stdn = fit_params.get(self._STDN, False)
        if stdn:
            self.apply_standard_scaler(X)

        # Use saved results
        use_cached = fit_params.get('use_cached', False)
        if use_cached and self.df_res is not None:
            df_res = self.df_res.copy()

        else:
            df_res = calculate_granger_causation(
                X, x_vars=self.x_vars, y_vars=self.y_vars,
                shift=self.shift, test=self.test, 
                lag_order=self.lag_order, max_lag=self.max_lag,    
                n_jobs=self.n_jobs
            )

        # Store results
        self.df_res = df_res.copy()

        log2 = fit_params.get(self._LOG2, False)
        if log2:
            df_res = apply_log2_fold(df_res)

        scor = fit_params.get(self._SCOR, False)
        if scor:
            df_res = apply_signed_correlation(X, df_res)        

        return df_res
    
    def fit_transform(self, X: DataFrame, y = None, **fit_params):
        df_res = self.fit(X).transform(X, y = y, **fit_params)
        return df_res
    
    def apply_log2_fold(self, df: DataFrame) -> DataFrame:
        return apply_log2_fold(df)

    def apply_signed_correlation(self, df_tseries: DataFrame, df_granger: DataFrame) -> pd.DataFrame:
        return apply_signed_correlation(df_tseries, df_granger)
    
    def apply_standard_scaler(self, df: DataFrame) -> DataFrame:
        df_res, scaler = apply_standard_scaler(df, return_scaler=True)
        self.scaler = scaler
        return df_res
    
    def invert_scaler(self, df: DataFrame) -> DataFrame:
        if self.scaler is None:
            raise ValueError
        return self.scaler.inverse_transform(df)
    
    def plot_df_org(self, show_all_yticks: bool = True, **kwargs):
        '''
        Returns
        -------
        fig : matplotlib.pyplot.Figure
        ax : matplotlib.pyplot.Axis
        ClusterMap : seaborn.ClusterMap
        '''
        options = dict(cmap='inferno', robust=True, col_cluster=False, yticklabels=show_all_yticks)
        options.update(kwargs)
        cstrmp = sns.clustermap(self.df_org, **options)
        return cstrmp
    
    def plot_df_res(self, show_all_yticks: bool = True, **kwargs):
        '''
        Returns
        -------
        fig : matplotlib.pyplot.Figure
        ax : matplotlib.pyplot.Axis
        ClusterMap : seaborn.ClusterMap
        '''           
        options = dict(yticklabels=show_all_yticks)
        options.update(kwargs)
        cstrmp = sns.clustermap(self.df_res, **options)
        return cstrmp

