import pandas as pd
import numpy as np


class SafeAI:
    """
    This class represents proposed functions in "A Rank Graduation Box for SAFE AI":

    Babaei, G., Giudici, P., & Raffinetti, E. (2024). A Rank Graduation Box for SAFE AI. 
    Expert Systems with Applications, 125239.

    It demonstrates how to measure Accuracy, Explainability, Fairness, 
    and Robustness of Machine Learning models based on the Lorenz curve.
    """
    def __init__(self):
        pass
    
    def rga(self, y, yhat):
        """
        RANK GRADUATION ACCURACY (RGA) MEASURE 
        Function for the RGA measure computation

        Inputs:
        y    : Actual values of the target variable 
        yhat : Predicted values generated by the model 
        
        Returns:
        RGA  : Calculated RGA measure
        """
        # Convert inputs to DataFrames and concatenate them
        y = pd.DataFrame(y).reset_index(drop=True)
        yhat = pd.DataFrame(yhat).reset_index(drop=True)

        df = pd.concat([y, yhat], axis=1)
        df.columns = ["y", "yhat"]
        
        # Rank yhat values
        df['ryhat'] = df['yhat'].rank(method="min")

        # Group by ryhat and calculate mean of y (support)
        support = df.groupby('ryhat')['y'].mean().reset_index(name='support')

        # Merge support back to the original dataframe
        df = pd.merge(df, support, on="ryhat", how="left")

        # Create the rord column by directly assigning 'support' where ryhat matches
        df['rord'] = df['support']
        
        # Sort df by yhat to get correct ordering for ystar
        df = df.sort_values(by="yhat").reset_index(drop=True)

        # Get ystar in the same order as rord in the sorted dataframe
        ystar = df['rord'].values

        # Create the index array I
        I = np.arange(len(df))

        # Calculate conc, dec (descending order of y) and inc (ascending order of y)
        conc = np.sum(I * ystar)
        sorted_y = np.sort(df['y'])  # y sorted in ascending order
        dec = np.sum(I * sorted_y[::-1])  # y sorted in descending order
        inc = np.sum(I * sorted_y)

        # Compute the RGA
        RGA = (conc - dec) / (inc - dec)

        return RGA

    def rge(self, yhat, yhat_rm):
        """
        RANK GRADUATION EXPLAINABILITY (RGE) MEASURE
        Function for the RGE measure computation

        Inputs:
        yhat    : Predicted values generated by the full model including all the explanatory variables 
        yhat_rm : Predicted values generated by the model excluding the selected variable
        
        Returns:
        RGE  : Calculated RGE measure
        """ 
        yhat = pd.DataFrame(yhat).reset_index(drop=True)
        yhat_rm = pd.DataFrame(yhat_rm).reset_index(drop=True)
        rge = 1-(self.rga(yhat, yhat_rm))
        return rge

    def rgf(self, yhat, yhat_rm):
        """
        RANK GRADUATION FAIRNESS (RGF) MEASURE 
        Function for the RGF measure computation

        Inputs:
        yhat    : Predicted values generated by the full model including all the explanatory variables 
        yhat_rm : Predicted values generated by the model excluding the selected variable
        
        Returns:
        RGF  : Calculated RGF measure
        """
        yhat = pd.DataFrame(yhat).reset_index(drop=True)
        yhat_rm = pd.DataFrame(yhat_rm).reset_index(drop=True)
        rgf = self.rga(yhat, yhat_rm)
        return rgf

    def rgr(self, yhat, yhat_pert):
        """
        RANK GRADUATION Robustness (RGR) MEASURE
        Function for the RGR measure computation regarding perturbation of a single variable

        Inputs:
        yhat    : Predicted values generated by the full model  
        yhat_pert : Predicted values generated by the model including the perturbed selected variable
        
        Returns:
        RGR  : Calculated RGR measure
        """ 
        yhat = pd.DataFrame(yhat).reset_index(drop=True)
        yhat_pert = pd.DataFrame(yhat_pert).reset_index(drop=True)
        rgr = self.rga(yhat, yhat_pert)
        return rgr