from pprint import pp
import matplotlib.pyplot as plt
import numpy as np
from aare import add_colorbar

def plot_histogram(histogram_data : np.ndarray, bin_edges : np.ndarray, Count_range : tuple[int, int] = None, bin_range : tuple[int, int] = None, label : str = None, title : str = None, xlabel : str = None, axis : plt.Axes = None) -> plt.Axes:
    """
    Plot the histogram of the pedestal values.

    Parameters
    ----------
    histogram_data : np.ndarray
        The histogram data to plot.
    bin_edges : np.ndarray
        The edges of the bins for the histogram.
    Count_range : tuple[int, int], optional
        The range of counts to display on the y-axis. Default is None.
    bin_range : tuple[int, int], optional
        The range of bin values to display on the x-axis. If None, it will be set to the range of the bin edges. Default is None.
    label : str, optional
        The label for the histogram. Default is None.
    title : str, optional
        The title for the plot. Default is None.
    xlabel : str, optional
        The label for the x-axis. Default is None.
    axis : plt.Axes, optional
        The axis to plot on. If None, a new figure and axis will be created.

    Returns
    -------
    plt.Axes
        The axis with the histogram plot.

    """

    if axis is None: 
        fig, ax = plt.subplots(figsize = (8,5))
    else:
        ax = axis

    ax.stairs(histogram_data, bin_edges, label = label, zorder = 3)

    if bin_range is None:
        bin_range = (bin_edges[0]-0.01*bin_edges[0], bin_edges[-1]+0.01*bin_edges[-1])

    if Count_range is not None:
        ax.set_ylim(*Count_range)

    ax.set_xlim(*bin_range)
    ax.grid(zorder = 0)
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel('Counts')
    ax.set_title(title)

    if label is not None:
        ax.legend()

    return ax 

def plot_fitted_function(histogram_data : np.ndarray, bin_edges : np.ndarray, function : callable, Count_range : tuple[int, int] = None, bin_range : tuple[int, int] = None, label : str = None, xlabel : str = None, title : str = None, axis : plt.Axes = None) -> plt.Axes:

    """
    Plot the histogram of the pedestal values along with the fitted function.

    Parameters
    ----------
    histogram_data : np.ndarray
        The histogram data to plot.
    bin_edges : np.ndarray
        The edges of the bins for the histogram.
    function : callable
        The fitted function to plot.
    Count_range : tuple[int, int], optional
        The range of counts to display on the y-axis. Default is None.
    bin_range : tuple[int, int], optional
        The range of bin values to display on the x-axis. If None, it will be set to the range of the bin edges. Default is None.
    label : str, optional
        The label for the fitted function. Default is None.
    xlabel : str, optional
        The label for the x-axis. Default is None.
    title : str, optional
        The title for the plot. Default is None.
    axis : plt.Axes, optional
        The axis to plot on. If None, a new figure and axis will be created.

    Returns
    -------

    plt.Axes
        The axis with the histogram and fitted function plot.
    
    """

    ax = plot_histogram(histogram_data, bin_edges, Count_range = Count_range, bin_range = bin_range, xlabel = xlabel, axis = axis, title = title)

    bin_centers = bin_edges[:-1] + np.diff(bin_edges)/2

    ax.plot(bin_centers, function(bin_centers), label=label, color='red', zorder=4)

    if label is not None:
        ax.legend(prop={"family": "monospace", "size": 10}, loc= 'upper left')

    return ax 

def plot_parameter(fit_parameter : np.ndarray, parameter_name : str, suppress_outliers : bool = True) -> None: 
    """
    Plot the parameter for all pixels.

    Parameters
    ----------
    fit_parameter : np.ndarray
        The parameter values to plot.
    parameter_name : str
        The name of the parameter to plot - used for plot title. 
    suppress_outliers : bool, optional
        Whether to suppress outliers in the plot (True, don't plot outliers). Default is True.
    """

    fig, ax = plt.subplots(figsize = (15,5))
    im = ax.imshow(fit_parameter)
    #suppress outliers for colorbar
    if suppress_outliers:
        mean = np.mean(fit_parameter)
        std = np.std(fit_parameter)
        im.set_clim(mean-3*std,mean+3*std)
    ax.set_xlabel('Pixel X')
    ax.set_ylabel('Pixel Y')
    ax.set_title(parameter_name)
    add_colorbar(ax, im)
    plt.show()