# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/11_DSCM_plots.ipynb.

# %% auto 0
__all__ = ['plot_time_series', 'plot_hist', 'plot_time_series_and_hist', 'plot_time_series_split', 'plot_decision_on_time_series',
           'PlotDecisionSliderHist', 'plot_lines_area', 'plot_lines_individual', 'PlotDecisionSliderMarg',
           'plot_cost_and_marginal_cost', 'plot_normal_dist']

# %% ../../nbs/11_DSCM_plots.ipynb 3
# Import necessary libraries
from .curly_brace import curlyBrace
from .cost import * 


from matplotlib import pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
from matplotlib.lines import Line2D

from ipywidgets import interact, FloatSlider, widgets
from IPython.display import clear_output
from scipy.stats import norm



def plot_time_series(demand, ax):
    """
    Plot the time series on the left subplot

    Parameters
    ----------
    demand : array-like
        Array of demand values
    ax : matplotlib.axes.Axes
        The axes object to plot on
    
    Returns
    -------
    None
    
    """
    
    ax.plot(demand, color="blue", linewidth=0.5)  # Set the color to blue
    ax.set_xlabel("Time", fontsize=12) # Define the labels for the x-axsis
    ax.set_ylabel("Demand", fontsize=12) # Define the labels for the y-axsis
    ax.set_title("Time Series of Demand", fontsize=12) # Set the title
    ax.tick_params(axis="both", labelsize=12) # Set the size of the ticks of the axes
    

def plot_hist(demand, ax):
    """
    Plot the histogram on the right subplot

    Parameters
    ----------
    demand : array-like
        Array of demand values
    ax : matplotlib.axes.Axes
        The axes object to plot on

    Returns
    -------
    None

    """
    ax.hist(demand, bins=40, color="blue")  # Set the color to blue
    ax.set_xlabel("Demand", fontsize=11) # Define the labels for the x-axsis
    ax.set_ylabel("Frequency", fontsize=11) # Define the labels for the y-axsis
    ax.set_title("Histogram of Demand", fontsize=11) # Set the title
    ax.tick_params(axis="both", labelsize=11) # Set the size of the ticks of the axes


def plot_time_series_and_hist(demand):
    """
    Plot the time series and histogram of the demand

    Parameters
    ----------
    demand : array-like
        Array of demand values

    Returns
    -------
    None

    """

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    plot_time_series(demand, ax1)

    plot_hist(demand, ax2)


def plot_time_series_split(demand, horizontal_line=None):
    """
    Plot the time series split into train and test set

    Parameters
    ----------

    demand : array-like
        Array of demand values
    horizontal_line : float, optional
        Value of the horizontal line to plot

    Returns
    -------
    None

    """

    # Create a figure with a single subplot
    fig, ax = plt.subplots(figsize=(15, 5))

    # Create an array of x-values
    x = np.arange(len(demand))
    
    # Plot the time series (train set) in blue
    ax.plot(x[:500], demand[:500], color="blue", linewidth=.5, label="Train set")  # Set the color to blue and add a label for the train set

    # Plot the line on the right side of the vertical line (test set) in grey
    ax.plot(x[500:], demand[500:], color="lightgrey", linewidth=.5, label="Test set")  # No label for the test set in this plot call

    # Add a vertical line at 500
    ax.axvline(x=500, color="black", linestyle="-", linewidth=1, label="Today")  # Add a label for the vertical line

    if horizontal_line is not None:
        ax.axhline(y=horizontal_line, color="red", linestyle="-", linewidth=5, label="Average demand")  # Add a label for the vertical line

    # Reorder the legend handles and labels to have "Today" at the beginning
    handles, labels = ax.get_legend_handles_labels()
    order = [labels.index("Today")] + [i for i in range(len(labels)) if labels[i] != "Today"]
    handles = [handles[i] for i in order]
    labels = [labels[i] for i in order]
     
    # Add a legend for train set, test set, and horizontal line
    ax.legend(handles, labels, fontsize=11)
    ax.set_xlabel("Time", fontsize=11) # Define the labels for the x-axsis
    ax.set_ylabel("Demand", fontsize=11) # Define the labels for the y-axsis
    ax.set_title("Time Series of Demand", fontsize=11) # Set the title
    ax.tick_params(axis="both", labelsize=11) # Set the size of the ticks of the axes

def plot_decision_on_time_series(demand, quantity):
    """
    Plot the decision on the time series of the demand which is split into train and test set

    Parameters
    ----------
    demand : array-like
        Array of demand values
    quantity : float
        Quantity to order

    Returns
    -------
    None

    """

    # Create an array of x-values
    x = np.arange(len(demand))

    # Create a figure with a single subplot
    fig, ax = plt.subplots(figsize=(15, 5))

    # Plot the time series (train set) in blue
    ax.plot(x[:500], demand[:500], color="blue", linewidth=.5, label="Train set")  # Set the color to blue and add a label for the train set

    # Generate the arrays for the test set
    test_set_max = np.maximum(quantity, demand[500:])
    test_set_min = np.minimum(quantity, demand[500:])

    # Plot the line on the right side of the vertical line (test set) in green (max values) and orange (min values)
    ax.plot(x[500:], test_set_max, color="green", linewidth=.5, label="Test set (high demand)")  # Add a label for the test set max values
    ax.plot(x[500:], test_set_min, color="orange", linewidth=.5, label="Test set (low demand)")  # Add a label for the test set min values

    # Add a vertical line at 500
    ax.axvline(x=500, color="black", linestyle="-", linewidth=1, label="Time of decision making")  # Add a label for the vertical line

    # Add a horizontal line at q
    ax.axhline(y=quantity, color="red", linestyle="-", linewidth=.5, label="Order quantity")

    # Reorder the legend handles and labels to have "Time of decision making" at the beginning
    handles, labels = ax.get_legend_handles_labels()
    order = [labels.index("Time of decision making")] + [i for i in range(len(labels)) if labels[i] != "Time of decision making"]
    handles = [handles[i] for i in order]
    labels = [labels[i] for i in order]

    # Add a legend for train set, test set max values, test set min values, and horizontal line
    ax.legend(handles, labels, fontsize=11)

    # Add a legend for train set, test set, and horizontal line
    ax.legend(handles, labels, fontsize=11)
    ax.set_xlabel("Time", fontsize=11) # Define the labels for the x-axsis
    ax.set_ylabel("Demand", fontsize=11) # Define the labels for the y-axsis
    ax.set_title("Time Series of Demand", fontsize=11) # Set the title
    ax.tick_params(axis="both", labelsize=11) # Set the size of the ticks of the axes

    # Extend the x-axis limit to 1000
    ax.set_xlim([0, 1000])

    # Add annotations for demand higher and lower than defrosted ingredients
    max_demand = np.max(demand[500:])
    min_demand = np.min(demand[500:])

    heights = [quantity - min_demand, max_demand - quantity]  # Heights of the brackets
    
    
    # curly bracket 1 start point and end point
    start_overage = [730, quantity]
    end_overage = [730, min_demand]
    
    # curly bracket 2 start point and end point
    start_underage = [730, max_demand]
    end_underage = [730, quantity]


    # fontdict for curly bracket 1 text
    font_overage = {'family': 'serif',
            'color':  'orange',
            'weight': 'bold',
            'style': 'italic',
            'size': 20,
            }
    
        # fontdict for curly bracket 1 text
    font_underage = font_overage.copy()
    font_underage["color"] = "green"

    # curly brack text
    str_text_overage = 'Demand higher'
    str_text_underage = 'Demand higher'

    # coefficient for curly bracket 1
    coefficients_bracket_overage = 0.08
    coefficients_bracket_underage = 0.05

    # clockwise, need to swap the start point and end point
    curlyBrace(fig, ax, start_overage, end_overage, coefficients_bracket_overage, bool_auto=True, str_text=None, color=font_overage["color"], lw=.5, int_line_num=1, fontdict=font_overage)

    curlyBrace(fig, ax, start_underage, end_underage, coefficients_bracket_underage, bool_auto=True, str_text=None, color=font_underage["color"], lw=.5, int_line_num=1, fontdict=font_underage)

    # add text description for the curly braces:
    y_overage = (max_demand + quantity) / 2
    y_underage = (min_demand + quantity) / 2
    
    ax.text(755, y_overage, 'Higher demand than prepared meals, \nunderage cost cu incurred per unit \nof demand not served', color='green', fontsize=10, verticalalignment='center', horizontalalignment='left')
    ax.text(755, y_underage, 'Lower demand than prepared meals, \noverage cost co incurred per \ndiscarded unit', color='orange', fontsize=10, verticalalignment='center', horizontalalignment='left')


    # # Display the figure
    # plt.show()


from matplotlib import gridspec

class PlotDecisionSliderHist():

    """
    Plot the demand time series and histogram with a slider to change the decision quantity

    Parameters
    ----------
    demand : array-like
        Array of demand values
    cu : float
        Underage cost
    co : float
        Overage cost
    
    Returns
    -------
    None

    """

    def __init__(self, demand, cu, co):
        self.demand = demand
        self.cu = cu
        self.co = co

        q = 0

        self.plot_decision_slider_hist_init()
        
    
        # Function to plot the figure
    def plot_figure_init(self, demand, q):
    
        
        self.fig = plt.figure(figsize=(18, 5)) 
        self.gs = gridspec.GridSpec(1, 3, width_ratios=[1, 1, 0.3])
        
        self.ax1 = plt.subplot(self.gs[0])  # First plot in the first column
        self.ax2 = plt.subplot(self.gs[1])
        
        plot_time_series(demand, self.ax1)
        plot_hist(demand, self.ax2)
    
        
        self.update_slider(q)

    def update_slider(self, q):
        # Draw red line on time series
        self.ax1.axhline(y=q, color='red', linewidth=.5)
        red_legend_artist = Line2D([0], [0], color='red', linewidth=.5, label='Quantity for de-frosting')
        self.ax1.legend(handles=[red_legend_artist], fontsize=11)
        
        # Draw red line on histogram
        self.ax2.axvline(x=q, color='red', linewidth=.5)
        red_legend_artist = Line2D([0], [0], color='red', linewidth=.5, label='Quantity for de-frosting')
        self.ax2.legend(handles=[red_legend_artist], fontsize=11)

        # Calculate cost
        total_avg, cu_avg, co_avg = get_exp_cost(y_true = self.demand, y_pred = [q]*len(self.demand), cu=self.cu, co=self.co, return_total_cu_co = True)
        total_avg, cu_avg, co_avg = np.round(total_avg, 5), np.round(cu_avg, 5), np.round(co_avg, 5)

        
        ax_text = plt.subplot(self.gs[2])
        ax_text.axis('off')  # Turn off axis
        label_width = max(len("Underage cost:"), len("Overage cost:"), len("Total cost:"))
        value_start = label_width + 5  # Add some space between labels and values
        num_underscores = 43

        # Prepare the text with aligned values
        textstr = (
            f"{'Average daily cost on train set:':{value_start}}\n"
            f" \n"
            f"{'-' * num_underscores}\n"
            f"{'Underage cost:':<{label_width}} {'':5} {cu_avg}\n"
            f"{'Overage cost:':<{label_width}} {'':6} {co_avg}\n"
            f"{'-' * num_underscores}\n"
            f"{'Total cost:':<{label_width}} {'':10} {total_avg}"
        )
        
        #textstr = f'Average daily cost on train set:\n - Underage cost: {cu_avg}\n - Overage cost: {co_avg}\n - Total cost: {total_avg}'
        ax_text.text(0.0, 0.5, textstr, transform=ax_text.transAxes, fontsize=14,
        verticalalignment='center', horizontalalignment='left')
    
        
        self.fig.canvas.draw()

    # Function to update the plot based on the slider value
    def update_plot_1(self, q):
        
        # clear_output(True)
        self.update_slider(q)

    def plot_decision_slider_hist_init(self):

        q_init = 0 # Initial value for q
        # Create the slider
        q_slider = FloatSlider(min=0, max=100, step=5, value=q_init, description='q')
        # Display the slider and the plot using interact
        interact(self.plot_figure_init, demand=widgets.fixed(self.demand), q=q_slider, continuous_update=False)
        

def plot_lines_area(total_costs, cu_costs, co_costs, q_tried, ax):
    
    total_costs = total_costs if isinstance(total_costs, np.ndarray) else np.array(total_costs) 
    cu_costs = cu_costs if isinstance(cu_costs, np.ndarray) else np.array(cu_costs) 
    co_costs = co_costs if isinstance(co_costs, np.ndarray) else np.array(co_costs) 
    
    ax.plot(q_tried, cu_costs, label='CU Cost', color = "orange")
    ax.plot(q_tried, co_costs + cu_costs, label='CO Cost', color = "green")
    ax.plot(q_tried, total_costs, label='Total Cost', color = "blue")

    ax.set_xlabel('q', fontsize=11)
    ax.set_ylabel('Cost', fontsize=11)
    ax.set_title('Average daily costs per q', fontsize=11)
    ax.legend(fontsize=11)
    ax.set_xlim(0, 100)
    ax.tick_params(axis="both", labelsize=11)

    # Fill area under CU curve with orange color
    ax.fill_between(q_tried, 0, cu_costs, color='orange', alpha=0.3)

    # Fill region between CU and Total costs with green color
    ax.fill_between(q_tried, cu_costs, total_costs, color='green', alpha=0.3)


def plot_lines_individual(total_costs, cu_costs, co_costs, q_tried, ax, label_1 = 'CU Cost', label_2='CO Cost'):

    ax.plot(q_tried, cu_costs, label=label_1, color="orange")
    ax.plot(q_tried, co_costs, label=label_2, color="green")
    
    if total_costs is not None:
        ax.plot(q_tried, total_costs, label='Total Cost', color="blue")
    
    ax.set_xlabel('q', fontsize=11)
    ax.set_ylabel('Cost', fontsize=11)
    ax.set_title('Average daily costs per q', fontsize=11)
    ax.legend(fontsize=11)
    ax.tick_params(axis="both", labelsize=11)
    ax.set_xlim(0, 100)

def plot_lines_area(total_costs, cu_costs, co_costs, q_tried, ax):
    
    total_costs = total_costs if isinstance(total_costs, np.ndarray) else np.array(total_costs) 
    cu_costs = cu_costs if isinstance(cu_costs, np.ndarray) else np.array(cu_costs) 
    co_costs = co_costs if isinstance(co_costs, np.ndarray) else np.array(co_costs) 
    
    ax.plot(q_tried, cu_costs, label='CU Cost', color = "orange")
    ax.plot(q_tried, co_costs + cu_costs, label='CO Cost', color = "green")
    ax.plot(q_tried, total_costs, label='Total Cost', color = "blue")

    ax.set_xlabel('q', fontsize=11)
    ax.set_ylabel('Cost', fontsize=11)
    ax.set_title('Average daily costs per q', fontsize=11)
    ax.legend(fontsize=11)
    ax.set_xlim(0, 100)
    ax.tick_params(axis="both", labelsize=11)

    # Fill area under CU curve with orange color
    ax.fill_between(q_tried, 0, cu_costs, color='orange', alpha=0.3)

    # Fill region between CU and Total costs with green color
    ax.fill_between(q_tried, cu_costs, total_costs, color='green', alpha=0.3)


def plot_lines_individual(total_costs, cu_costs, co_costs, q_tried, ax, label_1 = 'CU Cost', label_2='CO Cost'):

    ax.plot(q_tried, cu_costs, label=label_1, color="orange")
    ax.plot(q_tried, co_costs, label=label_2, color="green")
    
    if total_costs is not None:
        ax.plot(q_tried, total_costs, label='Total Cost', color="blue")
    
    ax.set_xlabel('q', fontsize=11)
    ax.set_ylabel('Cost', fontsize=11)
    ax.set_title('Average daily costs per q', fontsize=11)
    ax.legend(fontsize=11)
    ax.tick_params(axis="both", labelsize=11)
    ax.set_xlim(0, 100)


class PlotDecisionSliderMarg():
    """
    Plot the demand time series and marginal cost with a slider to change the decision quantity

    Parameters
    ----------
    demand : array-like
        Array of demand values
    cu : float
        Underage cost
    co : float
        Overage cost

    Returns
    -------
    None

    """

    def __init__(self, demand, cu, co):
        self.demand = demand
        self.cu = cu
        self.co = co
        
        self.pre_calc()

        self.plot_decision_slider_marg_init()
        
        
    
    def pre_calc(self):
        
        self.total_costs, self.cu_costs, self.co_costs = calculate_costs_per_q(self.demand, self.cu, self.co, None, upper_bound=100)
        
        
    def plot_figure_init(self, demand, q):
        
        self.fig = plt.figure(figsize=(18, 5)) 
        self.gs = gridspec.GridSpec(1, 3, width_ratios=[1, 1, 0.3])
        
        self.ax1 = plt.subplot(self.gs[0])  # First plot in the first column
        self.ax2 = plt.subplot(self.gs[1])
        
        plot_hist(demand, self.ax1)
    
        
        self.update_slider(q)
    
    def update_slider(self, q):

        # use function from session 1 to plot the time series and the demand
        
        q = int(q)
        
        plot_lines_area(self.total_costs[:q+1], self.cu_costs[:q+1], self.co_costs[:q+1], range(q+1), self.ax2)
        #plot_lines_area([], [], [], [], self.ax2)

        # Draw red line on histogram
        self.ax1.axvline(x=q, color='red', linewidth=2)
        red_legend_artist = Line2D([0], [0], color='red', linewidth=2, label='Quantity for de-frosting')
        self.ax1.legend(handles=[red_legend_artist], fontsize=11)

        # Draw red line on cost chart
        self.ax2.axvline(x=q, color='red', linewidth=2)
        red_legend_artist2 = Line2D([0], [0], color='red', linewidth=2, label='Quantity for de-frosting (cost chart)')

        # Retrieve the existing handles and labels from the cost chart's legend
        handles, labels = self.ax2.get_legend_handles_labels()

        # Add the red line to the legend
        red_legend_artist = Line2D([0], [0], color='red', linewidth=2, label='Quantity for de-frosting')
        handles.append(red_legend_artist)
        labels.append('Quantity for de-frosting')

        # Update the legend with the combined list
        self.ax2.legend(handles=handles, labels=labels, fontsize=11)

        # Calculate cost
        total_avg, cu_avg, co_avg = get_exp_cost(y_true = self.demand, y_pred = [q]*len(self.demand), cu=self.cu, co=self.co, return_total_cu_co = True)
        total_avg, cu_avg, co_avg = np.round(total_avg, 5), np.round(cu_avg, 5), np.round(co_avg, 5)

        ax_text = plt.subplot(self.gs[2])
        ax_text.axis('off')  # Turn off axis
        label_width = max(len("Underage cost:"), len("Overage cost:"), len("Total cost:"))
        value_start = label_width + 5  # Add some space between labels and values
        num_underscores = 43

        # Prepare the text with aligned values
        textstr = (
            f"{'Average daily cost on train set:':{value_start}}\n"
            f" \n"
            f"{'-' * num_underscores}\n"
            f"{'Underage cost:':<{label_width}} {'':5} {cu_avg}\n"
            f"{'Overage cost:':<{label_width}} {'':6} {co_avg}\n"
            f"{'-' * num_underscores}\n"
            f"{'Total cost:':<{label_width}} {'':10} {total_avg}"
        )
        
        #textstr = f'Average daily cost on train set:\n - Underage cost: {cu_avg}\n - Overage cost: {co_avg}\n - Total cost: {total_avg}'
        ax_text.text(0.0, 0.5, textstr, transform=ax_text.transAxes, fontsize=14,
        verticalalignment='center', horizontalalignment='left')
    

        # Function to update the plot based on the slider value
    def update_plot_2(demand, q):

        q_tried.append(q)

        total_costs, cu_costs, co_costs = calculate_costs_per_q(demand, cu, co, q_tried, upper_bound=100)
        
        self.update_slider(q)
        
    def plot_decision_slider_marg_init(self):
        
        q_init=0
        
        q_slider = FloatSlider(min=0, max=100, step=5, value=q_init, description='q')
        
        # Display the slider and the plot using interact
        interact(self.plot_figure_init, demand=widgets.fixed(self.demand), q=q_slider, continuous_update=False)

def plot_cost_and_marginal_cost(demand, cu, co):
    """
    Plot the cost and marginal cost

    Parameters
    ----------
    demand : array-like
        Array of demand values
    cu : float
        Underage cost
    co : float
        Overage cost

    Returns
    -------
    None

    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    q_tried = list(range(0,101))
    
    total_costs, cu_costs, co_costs = calculate_costs_per_q(demand, cu, co, q_tried=q_tried, upper_bound=100)
    marginal_underage, marginal_overage = calculate_marginal_costs_per_q(demand, cu, co, q_tried=q_tried, upper_bound=100)
    
    plot_lines_area(total_costs, cu_costs, co_costs, q_tried, ax1)
    plot_lines_individual(None, marginal_underage, marginal_overage, q_tried, ax2,  label_1 = 'Benefit of an additional unit (saved underage costs)', label_2='Cost of an additional unit (arising overage costs)')
    
    min_cost = min(total_costs)
    min_cost_quantity = q_tried[total_costs.index(min_cost)]
    
    ax2.set_title("Average marginal cost per q", fontsize=11)
    
    # Add a red vertical line at the minimum cost quantity on ax1
    ax1.axvline(x=min_cost_quantity, color='red', label='Optimal Quantity')

    # Add a red vertical line at the minimum cost quantity on ax2
    ax2.axvline(x=min_cost_quantity, color='red', label='Optimal Quantity')
    
    #ax2.axhline(y=0)
    
    ax1.legend(fontsize=11)
    ax2.legend(fontsize=11)
    

    # Display the figure
    plt.show()




def plot_normal_dist(demand):
    """
    Plot the demand and fitted normal distribution

    Parameters
    ----------
    demand : array-like
        Array of demand values

    Returns
    -------
    None

    """

    

    mean = np.mean(demand)
    std = np.std(demand)

    normal = norm(mean,std)
    normal_x = np.linspace(np.min(demand), np.max(demand), 100)
    normal_y = normal.pdf(normal_x)

    # plot demand and fitted normal distibution
    plt.figure(figsize=(15, 5))
    plt.hist(demand, bins = 40, density=True, color="blue")
    plt.plot(normal_x, normal_y, color = "red")
    plt.show()
        
