import numpy as np
from sklearn.feature_selection import SelectKBest, f_classif
from scipy.stats import pearsonr
from sklearn.feature_selection import r_regression


def imbalance_ratio(data, col_label):
    """
    Returns IR (float), the imbalance ratio of a dataset
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    counts = [count for _, count in value_counts.items()]
    count_max = max(counts) # number of datapoints in majority class
    count_min = min(counts) # number of datapoints in minority class
    IR = count_max / count_min
    return IR

def adjusted_IR(data, col_label):
    """
    Returns adj_IR (float), the adjusted imbalance ratio of a dataset
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    IR = imbalance_ratio(data, col_label)
    lambda0 = 1 # constant value used in the paper presenting adjusted IR
    df = data.drop(columns=[col_label])
    labels = data[col_label]
    corr = df.corrwith(labels, method=lambda x, y: pearsonr(x, y)[0]) # PCC between each feature and the labels
    pvalues =  df.corrwith(labels, method=lambda x, y: pearsonr(x, y)[1]) # associated p-values
    tmp = [i for i in range(len(corr)) if corr.iloc[i] != 0 and pvalues.iloc[i] < 0.05]
    # p is the number of discriminative features = nb of features for which the correlation between the feature 
    # and the label is non-zero with sig. lvl <0.05
    p = len(tmp)
    if p == 0:
        p = 1
    adj_IR = IR - lambda0 * np.log2(p)
    return adj_IR

def hellinger_distance(vec1, vec2):
    """
    Returns the hellinger (float) distance between two vectors (arrays: vec1, vec2)
    """
    return np.sqrt(np.sum((np.sqrt(vec1) - np.sqrt(vec2)) ** 2)) / np.sqrt(2)

def entropy_class_proportions(data, col_label):
    """
    Returns C1 (float), the entropy of class proportions of a dataset
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    Note: C1 corresponds to imbalance_factor with version="Shannon"
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    size = len(data) # total number of datapoints
    proba = [(count/size) for _, count in value_counts.items()] # proba of appearance of each class
    summation = np.sum([(pc * np.log2(pc)) for pc in proba])
    C1 = -(1/np.log2(c)) * summation
    return C1

def multiclass_imbalance_ratio(data, col_label):
    """
    Returns C2 (float), the multiclass Imbalance Ratio
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    l_n = [count for _, count in value_counts.items()] # list of number of points in each class
    size = len(data) # total number of datapoints
    summation = np.sum([nc/(size - nc) for nc in l_n])
    IR_m = ((c - 1)/c) * summation
    C2 = 1 - (1/IR_m)
    return C2

def imbalance_degree(data, col_label):
    """
    Returns ID (float), the imbalance degree of a dataset with Hellinger distance
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    counts = [count for _, count in value_counts.items()]
    size = len(data) # total number of datapoints
    p = np.array(counts) / size # empirical distribution
    if (len(p) < c):
        p = np.concatenate((p, [0] * (c-len(p))))
    b = [1/c] * c # balanced distribution
    m = 0
    for prob in (p-b):
        if prob < 0:
            m +=1
    p_m = [0] * m + [1/c] * (c-m-1) + [1-(c-m-1)/c] # distribution with m minority class at its worst (0 pts in min class)
    # Hellinger distance is recommended in the paper introducing imbalance degree metric:
    # Hellinger computes the distance from empirical distribution to equiprobability distribution.
    ID = (hellinger_distance(p,b) / hellinger_distance(p_m,b)) + (m-1)
    return ID

def likelihood_ratio_imbalance_degree(data, col_label):
    """
    Returns LRID (float), the likelihood ratio imbalance degree of a dataset
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    size = len(data) # total number of datapoints
    proba = [(count/size) for _, count in value_counts.items()] # list of frequency of appearance of each class in dataset  
    l_n = [count for _, count in value_counts.items()] # list of number of points in each class
    summation = np.sum([(nc * -np.log(c*pc)) for (nc, pc) in zip(l_n, proba)])
    lrid = -2 * summation
    return lrid

def imbalance_factor(data, col_label, version="Collision"):
    """
    Returns IF (float), the imbalance factor of a dataset
    Several versions exist (Shannon, Collision, and Min), the paper introducting the imbalance factor mentions 
    that Collision performs best on average, therefore it is the one implemented by default here.
    Shannon and Min can also be used if specified with the arg version. If another name than Shannon, Min, or 
    Collision is used, the three versions are aggregated with a mean.
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    size = len(data) # total number of datapoints
    proba = [(count/size) for _, count in value_counts.items()] # proba of appearance of each class
    if version == "Collision":
        # Collision
        summation = np.sum([pc ** 2 for pc in proba])
        IF = - np.log2(summation) / np.log2(c)
    elif version == "Shannon":
        # Shannon
        summation = np.sum([pc * np.log2(pc) for pc in proba])
        IF = - summation / np.log2(c)
    elif version == "Min": 
        # Min
        summation = np.log2(max(proba))
        IF = - summation / np.log2(c)
    else: # Aggregation (mean) of 3 previous ones 
        # Collision
        summation = np.sum([pc ** 2 for pc in proba])
        IF = - np.log2(summation) / np.log2(c)
        # Shannon
        summation = np.sum([pc * np.log2(pc) for pc in proba])
        IF = IF - summation / np.log2(c)
        # Min
        summation = np.log2(max(proba))
        IF = IF - summation / np.log2(c)
        IF = IF/3
    return IF

def status_of_imbalance(data, col_label, version='whole'):
    """
    Returns SIMBA (float), the status of imbalance of a dataset (our proposed metric)
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    @version: String, for ablation study, if 'whole', SIMBA is returned, else if nonorm, noftimp, noftred, 
    SIMBA is returned without normalization, feature importance, or feature redundancy, respectively.
    """
    c = len(np.unique(data[col_label])) # number of classes
    value_counts = data[col_label].value_counts() # number of datapoints in each class
    size = len(data) # total number of datapoints
    proba = [(count/size) for _, count in value_counts.items()] # proba of appearance of each class
    n_c = [count for _, count in value_counts.items()] # number of points in each class
    labels = data[col_label]
    df = data.drop(columns=[col_label])

    # check for redundant features
    if version != 'noftred':
        corr = df.corr(method="pearson").abs()
        corr = corr[corr > 0.9] # features are considered redundant if their |PCC| > 0.9
        corr = corr.to_numpy()
        corr = np.tril(corr, k=0) # lower triangle
        corr[np.isnan(corr)] = 0
        indices = np.argwhere(corr != 0)
        indices = [list(i) for i in indices if i[0] != i[1]] # list of pair of indices of features that are redundant 
        features_to_rm = []
        if len(indices) != 0:
            # Calculate feature importance to know which feature to remove when two are redundant
            kbest = SelectKBest(f_classif, k="all")
            kbest.fit(df, labels)
            feature_imp = kbest.scores_
            columns = list(df.columns)
            for pair in indices:
                index1 = pair[0]
                index2 = pair[1]
                imp1 = feature_imp[index1]
                imp2 = feature_imp[index2]
                feature1 = columns[index1]
                feature2 = columns[index2]
                if imp1 >= imp2:
                    if feature2 not in features_to_rm:
                        features_to_rm.append(feature2)
                else:
                    if feature1 not in features_to_rm:
                        features_to_rm.append(feature1)
    else:
        features_to_rm = []
    
    if version != 'noftimp':
        p_c = []
        for cl in range(c):  # for each class 
            tmp = labels.replace(cl, -1)
            tmp[tmp >= 0] = 1
            tmp = tmp.replace(-1, 0)  # OVA structure class cl labeled 0, other classes labeled 1
            pcc = r_regression(df.drop(columns=features_to_rm), tmp)  # PCC between features and class cl vs all others
            pcc[np.isinf(pcc)] = 0
            # p = sum of all pcc / nb of non-redundant features, 
            # i.e. avg intensity of discrimination for non-redundant features
            p = np.sum(abs(pcc)) / (len(df.columns) - len(features_to_rm))  
            p_c.append(p) # list of % of informative features for each class
    
    if version == 'whole' or version == 'noftred':
        summation = np.sum([(pc * -np.log(c*pc)) for pc in proba])
        SIMBA = -2 * summation / (np.mean(p_c))
    elif version == 'nonorm':
        summation = np.sum([(n * -np.log(c*(n/size))) for n in n_c])
        SIMBA = -2 * summation / (np.mean(p_c))
    elif version == 'noftimp':
        summation = np.sum([(pc * -np.log(c*pc)) for pc in proba])
        SIMBA = -2 * summation
    return SIMBA

def measure_imbalance(data, col_label, rounded=False):
    """
    Returns the results obtained with each imbalance measure: 
    IR, adjIR, ID, LRID, IF, C1, C2, and SIMBA (all are floats)
    @data: pd.DataFrame, containing all data (features and labels)
    @col_label: String, name of the labels column in the pd.DataFrame data 
    @rounded: (optional) boolean, default is False, if True, the metric results are rounded to 2 decimals
    """
    IR = imbalance_ratio(data, col_label)
    adjIR = adjusted_IR(data, col_label)
    C1 = entropy_class_proportions(data, col_label)
    C2 = multiclass_imbalance_ratio(data,col_label)
    ID = imbalance_degree(data, col_label)
    LRID = likelihood_ratio_imbalance_degree(data, col_label)
    IF = imbalance_factor(data, col_label)
    SIMBA = status_of_imbalance(data, col_label)
    if rounded:
        IR = round(IR, 2)
        adjIR = round(adjIR, 2)
        C1 = round(C1, 2)
        C2 = round(C2, 2)
        ID = round(ID, 2)
        LRID = round(LRID, 2)
        IF = round(IF, 2)
        SIMBA = round(SIMBA, 2)
    return IR, adjIR, C1, C2, ID, LRID, IF, SIMBA

