
import pandas as pd
import numpy as np
from sklearn.metrics import roc_auc_score, roc_curve

class ModelEvaluator:
    def evaluate(self, y_true, y_pred_proba):
        """
        Calculate AUC, KS.
        """
        auc = roc_auc_score(y_true, y_pred_proba)
        
        fpr, tpr, thresholds = roc_curve(y_true, y_pred_proba)
        ks = max(tpr - fpr)
        
        return {
            'auc': auc,
            'ks': ks
        }
        
    def calculate_psi(self, expected, actual, buckets=10):
        """
        Calculate PSI between two distributions.
        expected: training scores (or reference)
        actual: test scores (or current)
        """
        def scale_range(input, min, max):
            input += -(np.min(input))
            input /= np.max(input) / (max - min)
            input += min
            return input

        breakpoints = np.arange(0, buckets + 1) / (buckets) * 100
        breakpoints = np.percentile(expected, breakpoints)
        
        # Avoid duplicate breakpoints
        breakpoints = sorted(list(set(breakpoints)))
        
        # Calculate counts
        expected_percents = np.histogram(expected, breakpoints)[0] / len(expected)
        actual_percents = np.histogram(actual, breakpoints)[0] / len(actual)
        
        def sub_psi(e_perc, a_perc):
            if a_perc == 0: a_perc = 0.0001
            if e_perc == 0: e_perc = 0.0001
            
            value = (e_perc - a_perc) * np.log(e_perc / a_perc)
            return value
            
        psi_value = np.sum([sub_psi(expected_percents[i], actual_percents[i]) for i in range(len(expected_percents))])
        
        return psi_value
