import math
import scipy.stats as st
import numpy as np
import pandas as pd
from scipy.stats import shapiro
from scipy.stats import pearsonr
from scipy.stats import spearmanr
from scipy.stats import kendalltau
from tools import get_type_columns
from sklearn.feature_selection import f_regression, mutual_info_regression


# CALCULATE Z VALUE
def get_z(confidence_level:float)->float:
    """
    Calculate Z value for a given confidence level.
    
    confidence_level -- confidence level into percent. 
    return -- z value.
    """
    return st.norm.ppf(1-(1-confidence_level/100.)/2)


# CALCULATE THE SAMPLE SIZE
def sample_size(population_size:int, confidence_level:float, confidence_interval:float)->int:
    """
    Calculate the sample size using the Cochran’s Sample Size Formula.
    
    population_size -- the total population size.
    confidence_level -- the seleceted confidence level in percent. 
    confidence_interval -- the selected confidence interval in percent.
    return -- sample size with the correction for smaller population (no large).
    """
    Z = 0.0
    p = 0.5
    e = confidence_interval/100.0
    N = population_size
    n_0 = 0.0
    n = 0.0

    # FIND THE NUM STD DEVIATIONS FOR THAT CONFIDENCE LEVEL
    Z = get_z(confidence_level)

    if Z == 0.0:
        return -1

    # CALC SAMPLE SIZE
    n_0 = ((Z**2) * p * (1-p)) / (e**2)

    # ADJUST SAMPLE SIZE FOR FINITE POPULATION
    n = n_0 / (1 + ((n_0 - 1) / float(N)) )

    return int(math.ceil(n)) # THE SAMPLE SIZE


## Tests whether a data sample has a Gaussian distribution according Shapiro test
def test_shapiro(data:np.array, significance:float = 0.05, verbose:bool = False) ->bool:
    """
    Tests whether a data sample has a Gaussian distribution according Shapiro test.

    Parameters
    ----------
    data : np.array()
        Data to be tested.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information or not. The default is False.

    Returns
    -------
    bool
        If data has a Gaussian distribution or not.

    """
    # estimate cofindence level
    confidence = (1 - significance) * 100
    # test
    stat, p = shapiro(data)
    # display
    if verbose:
        print('stat=%.3f, p=%.3f' % (stat, p))
    # check result and return
    if p > significance:
        if verbose:
            print(f'Probably Gaussian (confidence level = {confidence}%)')
        return True
    else:
        if verbose:
            print(f'Probably not Gaussian (confidence level = {confidence}%)')
        return False
        
    
    
## Normality test analysis for several columns in a dataframe
def analysis_normality(df:pd.DataFrame, numerical_columns: list, significance:float = 0.05, verbose:bool = False)->(dict,list):
    """
    Normality test analysis for several columns in a dataframe.

    Parameters
    ----------
    df : pd.DataFrame
        Dataframe to be tested.
    numerical_columns : list
        Variables to be tested.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information. The default is False.

    Returns
    -------
    (dict,list)
        Results of test stored in a dictionary. List of reports.

    """

    # validate columns
    for c in numerical_columns:
        assert c in df.columns.tolist(), f'column "{c}" is required.'
    # estimate cofindence level
    confidence = (1 - significance) * 100 
    # initialize
    disnormal = dict()
    lreports = list()
    # loop of columns
    for col in numerical_columns:
        # collect data
        data = df[col].values
        # Shapiro test
        result_shapiro = test_shapiro(data, significance = significance, verbose = verbose)
        # add result
        disnormal[col] = result_shapiro
        # build report
        if result_shapiro:
            sreport = f'"{col}": Probably Gaussian (confidence level = {confidence}%) according to "Shapiro test"'
            # append
            lreports.append(sreport)
        else:
            sreport = f'"{col}": Probably NOT Gaussian (confidence level = {confidence}%) according to "Shapiro test"'
        # display
        if verbose:
            print(sreport)
    # return
    return disnormal, lreports
        
    
    
## Calculate Pearson's coefficient
def correlation_pearson(data1:np.array, data2:np.array, significance:float = 0.05, verbose:bool = False)->float:
    """
    Calculate Pearson's coefficient.

    Parameters
    ----------
    data1 : np.array
        First data array to be used.
    data2 : np.array
        Second data array to be used.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information. The default is False.

    Returns
    -------
    float
        Correlation value.

    """
    # estimate cofindence level
    confidence = (1 - significance) * 100
    # calculate Pearson's correlation
    corr, p = pearsonr(data1, data2)
    # display
    # check result and return
    if p < significance:
        if verbose:
            print("Pearson's correlation: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return corr
    else:
        if verbose:
            print("Pearson's correlation is not trusted: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return np.nan    
    
    
    
## Calculate a Spearman correlation coefficient
def correlation_spearman(data1:np.array, data2:np.array, significance:float = 0.05, verbose:bool = False)->float:
    """
    Calculate a Spearman correlation coefficient.

    Parameters
    ----------
    data1 : np.array
        First data array to be used.
    data2 : np.array
        Second data array to be used.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information. The default is False.

    Returns
    -------
    float
        Correlation value.

    """
    # estimate cofindence level
    confidence = (1 - significance) * 100
    # calculate Pearson's correlation
    corr, p = spearmanr(data1, data2)
    # display
    # check result and return
    if p < significance:
        if verbose:
            print("Spearman's correlation: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return corr
    else:
        if verbose:
            print("Spearman's correlation is not trusted: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return np.nan  
    
    
    
## Calculate Kendall’s tau, a correlation measure for ordinal data
def correlation_kendalltau(data1:np.array, data2:np.array, significance:float = 0.05, verbose:bool = False)->float:
    """
    Calculate Kendall’s tau, a correlation measure for ordinal data.

    Parameters
    ----------
    data1 : np.array
        First data array to be used.
    data2 : np.array
        Second data array to be used.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information. The default is False.

    Returns
    -------
    float
        Correlation value.

    """
    # estimate cofindence level
    confidence = (1 - significance) * 100
    # calculate Pearson's correlation
    corr, p = kendalltau(data1, data2)
    # display
    # check result and return
    if p < significance:
        if verbose:
            print("Kendall’s tau correlation: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return corr
    else:
        if verbose:
            print("Kendall’s tau correlation is not trusted: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return np.nan      
    
    
    
## Calculate a F-test regression (linear dependency)
def f_test_regression(data1:np.array, data2:np.array, significance:float = 0.05, verbose:bool = False)->float:
 
    
    # estimate cofindence level
    confidence = (1 - significance) * 100
    # calculate F-test
    corr, p = spearmanr(data1, data2)
    # display
    # check result and return
    if p < significance:
        if verbose:
            print("Spearman's correlation: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return corr
    else:
        if verbose:
            print("Spearman's correlation is not trusted: %.3f (confidence interval = %s %s)"%(corr, confidence, '%'))
        return np.nan  

    

## Correlation analysis for a couple of columns in a dataframe (for numerical and ordinal data)
def analysis_correlation(data:pd.DataFrame, couple_columns: list, dnormality:dict = None, significance:float = 0.05, verbose:bool = False)->float:
    """
    ## Correlation analysis for a couple of columns in a dataframe (for numerical and ordinal data).

    Parameters
    ----------
    data : pd.DataFrame
        Dataframe to be used.
    couple_columns : list
        Couple of columns names to be used.
    dnormality : dict, optional
        Dictionary with results of normality test per column . The default is None.
    significance : float, optional
        Level of significance. The default is 0.05.
    verbose : bool, optional
        Display information. The default is False.

    Returns
    -------
    float
        Correlation value.

    """
 
    # validate columns
    for c in couple_columns:
        assert c in data.columns.tolist(), f'column "{c}" is required.'
    # initialize
    col1 = couple_columns[0]
    col2 = couple_columns[1]
    # collect data and remove records with NaN values
    df = data[couple_columns].dropna()
    nsample = len(df)
    v1 = data[col1].values
    v2 = data[col2].values
    # validate type of data
    dtypecols = get_type_columns(df)
    typecol1 = dtypecols[col1]
    typecol2 = dtypecols[col2]
    # clean 
    del df

    ## analysis
    
    # case: ord - ord
    if typecol1 == 'ord' and typecol2 == 'ord':
        # get number of unique values of ordinal variables
        n_unique_var1 = len(list(set(list(v1))))
        n_unique_var2 = len(list(set(list(v2))))
        # select technique
        if n_unique_var1 >= 5 and n_unique_var2 >= 5:
            corr = correlation_spearman(v1, v2, significance = significance, verbose = verbose)
        else:
            corr = correlation_kendalltau(v1, v2, significance = significance, verbose = verbose)
    # case: num - num
    elif typecol1 == 'num' and typecol2 == 'num':
        # get normality
        if dnormality is None:
            is_norm_var1 = False
            is_norm_var2 = False
        else:
            is_norm_var1 = dnormality[col1]
            is_norm_var2 = dnormality[col2]    
        # select technique
        if nsample >= 100:
            corr = correlation_pearson(v1, v2, significance = significance, verbose = verbose)
        else:
            if is_norm_var1 and is_norm_var2:
                corr = correlation_pearson(v1, v2, significance = significance, verbose = verbose)
            else:
                corr = correlation_spearman(v1, v2, significance = significance, verbose = verbose)
    # case: num - ord
    elif typecol1 == 'num' and typecol2 == 'ord':
        # get number of unique values of ordinal variables
        n_unique_var2 = len(list(set(list(v2))))   
        # select technique
        if n_unique_var2 < 5:
            corr = correlation_kendalltau(v1, v2, significance = significance, verbose = verbose)
        else:
            if nsample >= 100:
                corr = correlation_pearson(v1, v2, significance = significance, verbose = verbose)
            else:
                corr = correlation_spearman(v1, v2, significance = significance, verbose = verbose)
    # case: ord - num
    elif typecol1 == 'ord' and typecol2 == 'num':
        # get number of unique values of ordinal variables
        n_unique_var1 = len(list(set(list(v1))))   
        # select technique
        if n_unique_var1 < 5:
            corr = correlation_kendalltau(v1, v2, significance = significance, verbose = verbose)
        else:
            if nsample >= 100:
                corr = correlation_pearson(v1, v2, significance = significance, verbose = verbose)
            else:
                corr = correlation_spearman(v1, v2, significance = significance, verbose = verbose)    
    # other case
    else:
        corr = np.na
        # display
        if verbose:
            print('It was not possible to identify correctly columns types.')      
    # return
    return corr
     
        
                
                
            
        
            
            
        
    
    
    