
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, feature_type='numeric'):
        """
        Calculate PSI between two distributions.
        expected: training values (or reference)
        actual: test values (or current)
        buckets: number of bins for numeric features
        feature_type: 'numeric' or 'categorical'
        """
        if feature_type == 'categorical':
            # For categorical features, calculate PSI directly from category distributions
            # Get unique categories
            all_categories = sorted(list(set(expected) | set(actual)))
            
            # Calculate distributions
            expected_counts = pd.Series(expected).value_counts(normalize=True)
            actual_counts = pd.Series(actual).value_counts(normalize=True)
            
            # Calculate PSI for each category
            psi_value = 0
            for cat in all_categories:
                e_perc = expected_counts.get(cat, 0.0001)
                a_perc = actual_counts.get(cat, 0.0001)
                
                # Avoid division by zero
                if e_perc > 0 and a_perc > 0:
                    psi_value += (e_perc - a_perc) * np.log(e_perc / a_perc)
            
            return psi_value
        else:
            # For numeric features, use binning approach
            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
