import numpy as np
from sklearn.calibration import calibration_curve
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

def analyze_models(y_test, X_test_predict, X_test_proba, model_name, bins=10, visualize=True):
    ### Classification report for model
    report = classification_report(y_test, X_test_predict)

    print('Classification Report for', model_name,':\n')
    # print(report)

    ### Compute calibration curve for model:
    # Access the predicted probabilities for the positive class
    X_test_proba_pos = X_test_proba[:, 1]
    prob_true, prob_pred = calibration_curve(y_test, X_test_proba_pos, n_bins=10)

    if visualize == True:
        # Plot reliability diagram
        fig, ax1 = plt.subplots(figsize=(10, 6))

        # Plot the calibration curve
        ax1.plot(prob_pred, prob_true, marker='o')
        ax1.plot([0, 1], [0, 1], linestyle='--', color='r')
        ax1.set_xlabel('Mean predicted probability')
        ax1.set_ylabel('Fraction of positives')
        ax1.set_title('Reliability Diagram for '+model_name)

        # Create a histogram to show the distribution of predicted probabilities
        ax2 = ax1.twinx()
        hist, bins, _ = ax2.hist(X_test_proba_pos, bins=bins, alpha=0.5)

        # Calculate bin centers
        bin_centers = 0.5 * (bins[:-1] + bins[1:])

        # Add count labels above the histogram bars
        for count, x in zip(hist, bin_centers):
            # Only put the label above non-zero bars
            if count > 0:
                ax2.text(x, count, str(int(count)), ha='center', va='bottom', color='grey')

        ax2.set_ylabel('Count')

        # Set the layout to prevent overlapping of plots
        fig.tight_layout()

        # Show the plot
        plt.show()

    return prob_true, prob_pred, fig


def expected_calibration_error(y_true, y_prob, n_bins=10):
    """
    Calculate the Expected Calibration Error (ECE).
    
    Parameters:
    y_true : array-like, shape (n_samples,)
        True binary outcomes (0 or 1).
    y_prob : array-like, shape (n_samples,)
        Predicted probabilities for the positive class.
    n_bins : int
        Number of bins to divide the predicted probabilities into.
    
    Returns:
    ece : float
        Expected Calibration Error.
    """
    # Convert inputs to numpy arrays
    y_true = np.array(y_true)
    y_prob = np.array(y_prob)
    
    # Create bins for predicted probabilities
    bin_edges = np.linspace(0, 1, n_bins + 1)
    bin_indices = np.digitize(y_prob, bin_edges, right=True) - 1
    
    # Initialize variables to compute ECE
    total_samples = len(y_true)
    ece = 0.0

    # Loop over each bin
    for i in range(n_bins):
        # Get the predicted probabilities and true outcomes for this bin
        bin_mask = bin_indices == i
        bin_size = np.sum(bin_mask)

        if bin_size > 0:
            # Calculate the average predicted probability for this bin
            avg_pred_prob = np.mean(y_prob[bin_mask])
            # Calculate the actual fraction of positives in this bin
            avg_true_prob = np.mean(y_true[bin_mask])
            # Calculate the contribution of this bin to ECE
            ece += (bin_size / total_samples) * np.abs(avg_pred_prob - avg_true_prob)
    
    return ece