import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import astropy.units as u
from scipy.stats import norm
#import ctapipe
from ctaplot.ana.ana import (
    energy_resolution_per_energy,
    angular_resolution_per_energy
)


def plot_errorbar(x_low, x_high, y, xlog=False, ax=None, **kwargs):
    """
    Parameters
    ----------
    x_low : array-like
        The lower bounds of the x values.
    x_high : array-like
        The upper bounds of the x values.
    y : array-like
        The y values.
    xlog : bool, optional
        If True, the x-axis will be logarithmic. Default is False.
    ax : matplotlib.axes.Axes, optional
        The axes on which to plot. If None, uses the current axes. Default is None.
    **kwargs : dict, optional
        Additional keyword arguments to pass to `ax.errorbar`.
    Returns
    -------
    ax : matplotlib.axes.Axes
        The axes with the plotted error bars.
    """
    ax = plt.gca() if ax is None else ax

    kwargs.setdefault("ls", "-")
    if xlog:
        x_mean = np.sqrt(x_low * x_high)
    else:
        x_mean = 0.5 * (x_low + x_high)
    x_err = ((x_high - x_mean), (x_mean - x_low))

    # Plot function
    ax.errorbar(
        x_mean,
        y,
        xerr=x_err,
        **kwargs,
    )

    if xlog:
        ax.set_xscale("log")
    return ax


def plot_fill_between(x_low, x_high, y_low, y_high, xlog=False, ax=None, alpha=0.3, **kwargs):
    """
    Plot a fill between function with optional transparency

    Parameters
    ----------
    x_low, x_high : array-like
        Lower and upper x values
    y_low, y_high : array-like
        Lower and upper y values
    xlog : bool
        If True, the x values are plotted in log scale
    ax : `~matplotlib.axes.Axes`
        Axes to plot on

    **kwargs : dict
        Keyword arguments passed to `~matplotlib.axes.Axes.fill_between`

    Returns
    -------
    ax : `~matplotlib.axes.Axes`
        Axes with the plot
    """
    ax = plt.gca() if ax is None else ax

    kwargs.setdefault("ls", "-")
    if xlog:
        x_mean = np.sqrt(x_low * x_high)
    else:
        x_mean = 0.5 * (x_low + x_high)

    ax.fill_between(
        x_mean,
        y_low,
        y_high,
        alpha=alpha,
        **kwargs,
    )

    if xlog:
        ax.set_xscale("log")
    return ax



def plot_energy_resolution(energy_low, energy_high, energy_resolution, energy_resolution_err=None, ax=None, alpha=0.3, **kwargs):
    """
    Plot the energy resolution as a function of the true energy

    Parameters
    ----------
    energy_low: `~astropy.units.Quantity`
        Lower bound of the true energy bin
    energy_high: `~astropy.units.Quantity`
        Upper bound of the true energy bin
    energy_resolution: `~astropy.units.Quantity`
        Energy resolution
    energy_resolution_err: (`~astropy.units.Quantity`, ~astropy.units.Quantity`)
        Energy resolution error (lower, upper)
    ax: `~matplotlib.axes.Axes`
        Axes to plot on
    alpha : float
        Transparency level for the fill between (0: fully transparent, 1: fully opaque)
    **kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.errorbar`

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes with the plot
    """

    ax = plot_errorbar(
        energy_low.to(u.TeV).value, energy_high.to(u.TeV).value, energy_resolution, xlog=True, ax=ax, **kwargs
    )

    if energy_resolution_err is not None:
        # Unpack the errors into lower and upper bounds
        err_lower, err_upper = energy_resolution_err
        ax = plot_fill_between(
            energy_low.to(u.TeV).value,
            energy_high.to(u.TeV).value,
            energy_resolution - err_lower,
            energy_resolution + err_upper,
            xlog=True,
            ax=ax,
            alpha=alpha,
            **kwargs,
        )

    # Style settings
    ax.set_xlabel(r"$E_\mathrm{True} / \mathrm{TeV}$")
    ax.set_ylabel("Energy resolution")
    ax.grid(True, which="both")

    return ax


energy_bins = np.logspace(np.log10(2.51e-02), 2, 19) * u.TeV


def energy_resolution(
    true_energy, reco_energy, percentile=68.27, confidence_level=0.95, bias_correction=False
):
    """
    Copied and reused from ctaplot
    Ref:https://github.com/cta-observatory/ctaplot/blob/d68dc609f5a0114a41ad8feaebc6ba9b9629ea25/ctaplot/ana/ana.py
    Cpute the true_energy resolution of true_energy as the percentile (68 as standard) containment radius of
    `true_energy-true_energy)/true_energy`
    with the lower and upper confidence limits defined by the given confidence level

    Parameters
    ----------
    true_energy: 1d numpy array of simulated energy
    reco_energy: 1d numpy array of reconstructed energy
    percentile: float
        <= 100

    Returns
    -------
    `numpy.array` - [energy_resolution, lower_confidence_limit, upper_confidence_limit]

    """
    assert len(true_energy) == len(reco_energy), "Les longueurs des tableaux doivent être égales."

    res = np.abs((reco_energy - true_energy) / true_energy)

    median_bias = np.median(reco_energy - true_energy)
    if bias_correction:
        reco_energy_corrected = reco_energy - median_bias
    else:
        reco_energy_corrected = reco_energy

    # Compute percentile using scaling s1
    resolution_value = np.percentile(np.abs((reco_energy_corrected - true_energy) / true_energy), percentile)
    lower_conf, upper_conf = percentile_confidence_interval(
        res, percentile=percentile, confidence_level=confidence_level
    )

    return np.array([resolution_value, lower_conf, upper_conf])


def merge_stats(energy_resolutions):
    """
    Merge the statistics tables of different runs.

    Parameters
    ----------
    energy_resolutions : array-like
        A 2D array where each row represents a different run and each column represents a different measurement.

    Returns
    -------
    numpy.ndarray
        A 2D array containing the following statistics along the first axis:
        - Mean
        - Standard deviation
        - Minimum value
        - Maximum value
        - Median
        - 16th percentile
        - 84th percentile
    
    """
    mean = np.mean(energy_resolutions, axis=0)
    std = np.std(energy_resolutions, axis=0)
    min_val = np.min(energy_resolutions, axis=0)
    max_val = np.max(energy_resolutions, axis=0)
    median = np.median(energy_resolutions, axis=0)
    percentile_16 = np.percentile(energy_resolutions, 16, axis=0)
    percentile_84 = np.percentile(energy_resolutions, 84, axis=0)
    return np.array([mean, std, min_val, max_val, median, percentile_16, percentile_84])


def percentile_confidence_interval(x, percentile=68, confidence_level=0.95):
    """

    Return the confidence interval for the qth percentile of x for a given confidence level

    REF:
    http://people.stat.sfu.ca/~cschwarz/Stat-650/Notes/PDF/ChapterPercentiles.pdf
    S. Chakraborti and J. Li, Confidence Interval Estimation of a Normal Percentile, doi:10.1198/000313007X244457

    Parameters
    ----------
    x: `numpy.ndarray`
    percentile: `float`
        0 < percentile < 100
    confidence_level: `float`
        0 < confidence level (by default 95%) < 1

    Returns
    -------
    """
    sorted_x = np.sort(x)
    z = norm.ppf(confidence_level)
    q = percentile / 100.0
    lower_index = max(0, int(len(x) * q - z * np.sqrt(len(x) * q * (1 - q))))
    upper_index = min(int(len(x) * q + z * np.sqrt(len(x) * q * (1 - q))), len(x) - 1)
    return sorted_x[lower_index], sorted_x[upper_index]


def compute_energy_resolutions(dataframes, energy_bins):
    """
    Compute the energy resolution for each run and return the results.

    Parameters
    ----------
    dataframes : list of pandas.DataFrame
        List of dataframes, each containing 'true_energy' and 'reco_energy' columns.
    energy_bins : array-like
        Bins to use for computing the energy resolution.

    Returns
    -------
    list of numpy.ndarray
        List of energy resolution arrays, one for each dataframe.
    """
    energy_resolutions = []

    for df in dataframes:
        true_energy = df["true_energy"].values * u.TeV
        reco_energy = df["reco_energy"].values * u.TeV
        _, res = energy_resolution_per_energy(true_energy, reco_energy, bins=energy_bins)
        energy_resolutions.append(res)

    return energy_resolutions


def load_and_preprocess(file_path):
    """
    Load and preprocess data from an HDF5 file.
    Parameters
    ----------
    file_path : str
        The path to the HDF5 file to be loaded.
    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the preprocessed data with renamed columns.
    """
  
    # Load the HDF5 file
    #df = pd.read_hdf(file_path, key="/dl2/event/subarray/Graph")
    df = pd.read_hdf(file_path)

    df.rename(columns={"Graph_reco_energy": "reco_energy", "Graph_reco_gamaness": "Graph_reco_gammaness"}, inplace=True)

    return df


def combine_energy_resolutions(file_paths, ax=None, **kwargs):
    """
    Combine the energy resolutions from multiple files and plot it

    Parameters
    ----------
    file_paths: list of str
        Paths to the data files
    ax: `~matplotlib.axes.Axes`, optional
        Matplotlib axes object to plot on
    **kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.errorbar`

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes with the plot
    """
    ax = plt.gca() if ax is None else ax
    
    dfs = [load_and_preprocess(file_path) for file_path in file_paths]
    energy_resolutions = compute_energy_resolutions(dfs, energy_bins)

    # Merge stats
    avg_resolution = merge_stats(energy_resolutions)

    energy_low = energy_bins[:-1]
    energy_high = energy_bins[1:]

    # Extract average resolution and confidence intervals
    resolution = avg_resolution[0][:, 0]  # Average resolution
    resolution_min = avg_resolution[2][:, 0]  # Minimum resolution
    resolution_max = avg_resolution[3][:, 0]  # Maximum resolution
    resolution_err = (resolution - resolution_min, resolution_max - resolution)

    ax = plot_energy_resolution(energy_low, energy_high, resolution, energy_resolution_err=resolution_err, ax=ax, alpha=0.3, **kwargs)

    df = dfs[0]
    if 'RandomForest_reco_energy' in df.columns and df['RandomForest_reco_energy'].notnull().all():
        true_energy = df["true_energy"].values * u.TeV
        reco_energy = df["RandomForest_reco_energy"].values * u.TeV
        _, res_rf = energy_resolution_per_energy(true_energy, reco_energy, bins=energy_bins)
        kwargs.setdefault("color", "black")
        kwargs["label"] = "Random Forest"
        ax = plot_energy_resolution(energy_low, energy_high, res_rf[:,0], ax=ax, **kwargs)
        ax.legend()
    return ax


def plot_energy_bias(energy_low, energy_high, energy_bias, energy_bias_err=None, ax=None, **kwargs):
    """ 
    Plot the energy bias as a function of the true energy

    Parameters
    ----------
    energy_low: `~astropy.units.Quantity`
        Lower bound of the true energy bin
    energy_high: `~astropy.units.Quantity`
        Upper bound of the true energy bin
    energy_bias: `~astropy.units.Quantity`
        Energy bias
    energy_bias_err: (`~astropy.units.Quantity`, `~astropy.units.Quantity`)
        Energy bias error (lower, upper)
    ax: `~matplotlib.axes.Axes`
        Axes to plot on
    **kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.errorbar`

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes with the plot
    """

    ax = plot_errorbar(
        energy_low.to(u.TeV).value, energy_high.to(u.TeV).value, energy_bias, xlog=True, ax=ax, **kwargs
    )
    if energy_bias_err is not None:
        ax = plot_fill_between(
            energy_low.to(u.TeV).value,
            energy_high.to(u.TeV).value,
            energy_bias-energy_bias_err[0],
            energy_bias+energy_bias_err[1],
            xlog=True,
            ax=ax,
            **kwargs,
        )

    # Style settings
    ax.set_xlabel(r"$E_\mathrm{True}$ / $\mathrm{TeV}$")
    ax.set_ylabel("Energy bias")
    ax.grid(True, which="both")

    return ax


def energy_bias(true_energy, reco_energy, bins=None):
    """
    Compute the relative energy bias per energy bin.

    Parameters
    ----------
    true_energy: `astropy.Quantity` (1d array)
        Simulated energies.
    reco_energy: `astropy.Quantity` (1d array)
        Reconstructed energies.
    bins: `astropy.Quantity` (1d array)
        Energy bins - if None, a default binning is used.

    Returns
    -------
    (energy_bins, bias): (astropy.Quantity, numpy.array)
        Energy bins and corresponding energy biases.
    """
    if bins is None:
        bins = np.logspace(np.log10(2.51e-02), 2, 19) * u.TeV  

    bias_e = []
    for i in range(len(bins) - 1):
        mask = (true_energy >= bins[i]) & (true_energy < bins[i + 1])
        if np.sum(mask) > 0:
            bias_e.append(np.mean((reco_energy[mask] - true_energy[mask]) / true_energy[mask]))
        else:
            bias_e.append(np.nan)

    return bins, np.array(bias_e)


def combine_mean_biases(file_paths, ax=None, **kwargs):
    """
    Combine the biases from multiple files and plot the mean bias.

    Parameters
    ----------
    file_paths: list of str
        Paths to the data files.
    ax: `~matplotlib.axes.Axes`, optional
        Matplotlib axes object to plot on.
    **kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.plot`

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes with the plot.
    """

    ax = plt.gca() if ax is None else ax
    
    dfs = [load_and_preprocess(file_path) for file_path in file_paths]
    
    biases = []
    for df in dfs:
        true_energy = df['true_energy'].values * u.TeV
        reco_energy = df['reco_energy'].values * u.TeV
        _, bias = energy_bias(true_energy, reco_energy, bins=energy_bins)
        biases.append(bias)
    
    
    # avc
    avg_bias = merge_stats(biases)
    
    energy_low = energy_bins[:-1]
    energy_high = energy_bins[1:]
    

    bias_mean = avg_bias[0]  # Mean bias
    bias_min = avg_bias[2]   # Minimum bias
    bias_max = avg_bias[3]   # Maximum bias
    bias_err = (bias_mean - bias_min, bias_max - bias_mean)
    
    plot_energy_bias(energy_low, energy_high, bias_mean, energy_bias_err=bias_err, ax=ax, **kwargs)

    df = dfs[0]
    if 'RandomForest_reco_energy' in df.columns:
        true_energy = df["true_energy"].values * u.TeV
        reco_energy = df["RandomForest_reco_energy"].values * u.TeV
        _, bias_rf = energy_bias(true_energy, reco_energy, bins=energy_bins)
        kwargs.setdefault("color", "black")
        ax = plot_energy_bias(energy_low, energy_high, bias_rf, ax=ax, label='Random Forest', **kwargs)
        ax.legend()
    
    return ax



def compute_angular_resolutions(dataframes, energy_bins):
    """
    Compute the angular resolution for each run and return the results.

    Parameters
    ----------
    dataframes : list of pandas.DataFrame
        List of dataframes containing the true and reconstructed energy and angles.
    energy_bins : array-like
        Bins to use for the energy resolution computation.

    Returns
    -------
    resolutions : list of numpy.ndarray
        List of angular resolution arrays for each dataframe.
    """
    resolutions = []

    for df in dataframes:
        true_energy = df["true_energy"].values * u.TeV
        true_alt = df["true_alt"].values * u.deg
        true_az = df["true_az"].values * u.deg
        reco_alt = df["Graph_reco_alt"].values * u.deg
        reco_az = df["Graph_reco_az"].values * u.deg
        _, res = angular_resolution_per_energy(true_alt, reco_alt, true_az, reco_az, true_energy, bins=energy_bins)
        resolutions.append(res)

    return resolutions


def plot_angular_resolution(energy_low, energy_high, angular_resolution, angular_resolution_err=None, ax=None, **kwargs):
    """"
    Plot the angular resolution as a function of energy

    Parameters
    ----------
    energy_low: `~astropy.units.Quantity`
        Lower energy bin edges
    energy_high: `~astropy.units.Quantity`
        Upper energy bin edges
    angular_resolution: `~astropy.units.Quantity`
        Angular resolution
    ax: `~matplotlib.axes.Axes`
        Axes
    kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.errorbar`

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes
    """

    ax = plot_errorbar(energy_low.to(u.TeV).value,
                       energy_high.to(u.TeV).value,
                       angular_resolution.to(u.deg).value,
                       xlog=True, ax=ax, **kwargs)

    if angular_resolution_err is not None:
        ax = plot_fill_between(energy_low.to(u.TeV).value,
                               energy_high.to(u.TeV).value,
                               angular_resolution.to(u.deg).value-angular_resolution_err[0].to(u.deg).value,
                               angular_resolution_err[1].to(u.deg).value+angular_resolution.to(u.deg).value,
                               xlog=True, ax=ax, **kwargs)

    ax.set_xscale("log")
    ax.set_xlabel(r"$E_\mathrm{True}$ / TeV")
    ax.set_ylabel("Angular Resolution / deg")
    ax.grid(True, which="both")

    return ax


def combine_angular_resolutions(file_paths, ax=None, **kwargs):
    """
    Combine angular resolutions from multiple files and plot the results.
    Parameters
    ----------
    file_paths : list of str
        List of file paths to the data files to be processed.
    ax : matplotlib.axes.Axes, optional
        Matplotlib Axes object to plot on. If None, a new Axes object is created.

    Returns
    -------
    matplotlib.axes.Axes
        The Axes object with the plotted angular resolutions.
    Notes
    -----
    This function processes multiple data files to compute and plot the average angular resolution
    and its confidence intervals. If the data contains RandomForest reconstructed altitude and azimuth,
    it also plots the angular resolution for the RandomForest reconstruction.
    """


    ax = plt.gca() if ax is None else ax
    
    dfs = [load_and_preprocess(file_path) for file_path in file_paths]
    resolutions = compute_angular_resolutions(dfs, energy_bins)

    # Merge stats
    avg_resolution = merge_stats(resolutions) * resolutions[0].unit

    energy_low = energy_bins[:-1]
    energy_high = energy_bins[1:]

    # Extract average resolution and confidence intervals
    resolution = avg_resolution[0][:, 0]  # Average resolution
    resolution_min = avg_resolution[2][:, 0]  # Minimum resolution
    resolution_max = avg_resolution[3][:, 0]  # Maximum resolution
    resolution_err = (resolution - resolution_min, resolution_max - resolution)

    kwargs.setdefault("alpha", 0.3)
    ax = plot_angular_resolution(energy_low, energy_high, resolution, angular_resolution_err=resolution_err, ax=ax, **kwargs)

    df = dfs[0]
    if 'RandomForest_reco_alt' in df.columns and 'RandomForest_reco_az' in df.columns and df['RandomForest_reco_alt'].notnull().all() and df['RandomForest_reco_az'].notnull().all():
        true_energy = df["true_energy"].values * u.TeV
        true_alt = df["true_alt"].values * u.deg
        true_az = df["true_az"].values * u.deg
        reco_alt = df["RandomForest_reco_alt"].values * u.deg
        reco_az = df["RandomForest_reco_az"].values * u.deg
        _, res_rf = angular_resolution_per_energy(true_alt, reco_alt, true_az, reco_az, true_energy, bins=energy_bins)
        kwargs.setdefault("color", "black")
        kwargs["alpha"] = 1
        kwargs["label"] = "Random Forest"
        ax = plot_angular_resolution(energy_low, energy_high, res_rf[:,0], ax=ax, **kwargs)
        ax.legend()
    return ax


def compute_auc(dataframes, energy_bins):
    """
    Compute the AUC for each run and return the results.

    Parameters
    ----------
    dataframes : list of pandas.DataFrame
        List of dataframes, each containing the data for a single experiment. Each dataframe must have the columns:
        - "true_energy": The true energy of the event.
        - "true_shower_id": The true type of the shower.
        - "Graph_reco_gammaness": The reconstructed gammaness of the event.
    energy_bins : array-like
        The energy bins to use for computing the AUC.

    Returns
    -------
    aucs : list of float
        List of AUC values, one for each dataframe.
    
    Compute the AUC for each run and return the results
    """
    aucs = []

    for df in dataframes:
        true_energy = df["true_energy"].values * u.TeV
        true_type = df["true_shower_id"].values
        gammaness = df["Graph_reco_gammaness"].values
        _, auc = roc_auc_per_energy(true_type, gammaness, true_energy, energy_bins=energy_bins)
        aucs.append(auc)

    return aucs


def combine_auc(file_paths, ax=None, **kwargs):
    """
    Combine the AUC from multiple files and plot it

    Parameters
    ----------
    file_paths: list of str
        Paths to the data files
    ax: `~matplotlib.axes.Axes`, optional
        Matplotlib axes object to plot on

    Returns
    -------
    ax: `~matplotlib.axes.Axes`
        Axes with the plot
    """

    ax = plt.gca() if ax is None else ax

    dfs = [load_and_preprocess(file_path) for file_path in file_paths]
    aucs = compute_auc(dfs, energy_bins)

    # Merge stats
    avg_auc = merge_stats(aucs)

    energy_low = energy_bins[:-1]
    energy_high = energy_bins[1:]

    # Extract average AUC and confidence intervals
    auc = avg_auc[0]  # Average AUC
    auc_min = avg_auc[2]  # Minimum AUC
    auc_max = avg_auc[3] # Maximum AUC
    auc_err = (auc - auc_min, auc_max - auc)

    ax = plot_auc_per_energy(energy_low, energy_high, auc, auc_err=auc_err, ax=ax, **kwargs)

    df = dfs[0]
    if 'RandomForest_class_reco' in df.columns and df['RandomForest_class_reco'].notnull().all():
        _, auc_rf = roc_auc_per_energy(df["true_shower_id"].values, df['RandomForest_class_reco'].values, df['true_energy'].values * u.TeV,  energy_bins=energy_bins)
        kwargs.setdefault("color", "black")
        kwargs["label"] = "Random Forest"
        ax = plot_auc_per_energy(energy_low, energy_high, auc_rf, ax=ax, **kwargs)
        ax.legend()
    return ax


def plot_auc_per_energy(energy_low, energy_high, auc, auc_err=None, ax=None, **kwargs):
    """
    Plot the AUC as a function of the energy

    Parameters
    ----------
    energy_low: ~astropy.units.Quantity
        lower energy bin edges
    energy_high: ~astropy.units.Quantity
        upper energy bin edges
    auc: array
        AUC values
    ax: ~matplotlib.axes.Axes
        Axes to plot on
    kwargs: dict
        Keyword arguments passed to `~matplotlib.pyplot.errorbar`

    Returns
    -------
    ax: ~matplotlib.axes.Axes
    """

    ax = plot_errorbar(energy_low.to(u.TeV).value, energy_high.to(u.TeV).value, auc, xlog=True, ax=ax, **kwargs)

    if auc_err is not None:
        ax = plot_fill_between(
            energy_low.to(u.TeV).value,
            energy_high.to(u.TeV).value,
            auc - auc_err[0],
            auc + auc_err[1],
            xlog=True,
            ax=ax,
            **kwargs,
        )

    ax.grid(True, which='both')
    ax.set_xlabel(r"$E_\mathrm{True} / \mathrm{TeV}$")
    ax.set_ylabel('AUC (gamma vs protons)')
    return ax


def roc_auc_per_energy(true_type, gammaness, true_energy, energy_bins=None, gamma_label=0, **roc_auc_score_opt):
    """
    Compute AUC score as a function of the true gamma energy.
    The AUC score is calculated in a gamma versus all fashion.

    Parameters
    ----------
    true_type: `numpy.ndarray`
        labels
    gammaness: `numpy.ndarray`
        likeliness of a particle to be a gamma
    true_energy: `numpy.ndarray`
        particles true energy
    energy_bins: `astropy.Quantity`
    gamma_label: label of gammas in `true_type` array
    roc_auc_score_opt: see `sklearn.metrics.roc_auc_score` options

    Returns
    -------
    energy_bins, auc_scores: `numpy.ndarray, numpy.ndarray`
    """
    from sklearn.metrics import roc_auc_score
    energy_bins = np.logspace(-2, 2, 10) * u.TeV if energy_bins is None else energy_bins

    binarized_label = (true_type == gamma_label).astype(int)

    auc_scores = []
    for i in range(len(energy_bins) - 1):
        gamma_mask = (true_type == gamma_label) & (true_energy >= energy_bins[i]) & (true_energy < energy_bins[i + 1])
        cosmic_mask = (true_type != gamma_label)
        mask = gamma_mask | cosmic_mask

        if np.count_nonzero(gamma_mask) > 0 and np.count_nonzero(cosmic_mask) > 0:  # be sure to have both gamma and cosmic to be able to compute an auc
            auc_score = roc_auc_score(binarized_label[mask], gammaness[mask], **roc_auc_score_opt)
            if auc_score < 0.5:
                auc_score = 1 - auc_score
            auc_scores.append(auc_score)
        else:
            auc_scores.append(np.nan)

    return energy_bins, np.array(auc_scores)
