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

# %% auto 0
__all__ = ['lr', 'cu', 'co', 'lrnw', 'weekdays', 'service_levels', 'service_level', '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', 'plot_time_series_base', 'plot_time_series_decomp',
           'plot_demand_per_day', 'plot_demand_per_month', 'plot_prescriptions_weekday_order_quantities',
           'plot_prescriptions_weekday', 'interactive_weekday_plot', 'plot_lrnw_core', 'plot_lin_reg',
           'plot_loss_functions']

# %% ../../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')
    
        
    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()
        

# %% ../../nbs/11_DSCM_plots.ipynb 4
from matplotlib import pyplot as plt
import seaborn as sns
import pandas as pd
from sklearn.linear_model import QuantileRegressor
import ipywidgets as widgets
import numpy as np  
from sklearn.linear_model import LinearRegression
from ..newsvendor import LinearRegressionNewsvendor

lr = LinearRegression()
cu = 15
co = 5
lrnw = LinearRegressionNewsvendor(cu, co)


weekdays = ['MON', 'TUE', 'WED', 'THU', 'FRI', 'SAT', 'SUN']

def plot_time_series_base(demand, ax=None, title=None):
    """
    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
    
    """

    if ax is None:
        fig, ax = plt.subplots(figsize=(15, 5))

    if title is None:
        title = "Time Series of Demand"
        
    
    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(title, fontsize=12) # Set the title
    ax.tick_params(axis="both", labelsize=12) # Set the size of the ticks of the axes
    
    return ax 
    



def plot_time_series_decomp(demand, month,  ax=None, title=None, trend=None, seasonality=None, prediction=None):

    ax = plot_time_series_base(demand, ax, title)

    # Show the month on the x-axis (every 30th entry of the month array)
    ax.set_xticks(range(0, len(demand), 30))
    ax.set_xticklabels(month[::30], rotation=45)

    # If trend is provided, plot it
    if trend is not None:
        slope, interc = trend
        x = np.arange(len(demand))
        ax.plot(x, slope * x + interc, color="orange", linestyle="--", label="Trend")
        ax.legend()

    # If seasonality is provided, plot it (seasonality shows the average demand per month)
    if seasonality is not None:
        ax.plot(seasonality, color="red", linestyle="--", label="Seasonality")
        ax.legend()


    if prediction is not None:
        ax.plot(prediction, color="green", linestyle="--", label="Prediction")
        ax.legend()


    return ax




def plot_demand_per_day(features, demand, train, days, q_SAA=None):
    """
    Create a stripplot of the demand per day of the week with a horizontal line indicating the SAA quantity.

    Parameters:
    features (DataFrame): DataFrame containing feature data.
    demand (Series): Series containing demand data.
    train (array-like): Indices or boolean mask for training data.
    days (list): List of days of the week in the desired order.
    q_SAA (float, Series, or None): SAA quantity. If None, no line is added.
    Returns:
    None
    """
    # Set up the figure and axis
    plt.figure(figsize=(10, 6))
    
    # Create the stripplot
    ax = sns.stripplot(x=features['weekday'][train], y=demand[train], jitter=True, dodge=True, order=days)
    
    # Add a horizontal line for the SAA quantity
    if q_SAA is not None:

        # Check if this a single value or an array / pandas Series
        if isinstance(q_SAA, (int, float)):
            plt.axhline(y=q_SAA, color='orange', linestyle='-', label="SAA", alpha=0.5)
        else:
            for i, day in enumerate(days):
                q_day = q_SAA.loc[day]
                plt.axhline(y=q_day, color='orange', linestyle='-', alpha=0.5, xmin=(i/7 + 0.015), xmax=((i+1)/7) - 0.015)

            plt.plot([], [], color='orange', linestyle='-', alpha=0.5, label="SAA per day")
    
        # Add the legend
        plt.legend()
    
    # Display the plot
    plt.show()


def plot_demand_per_month(features, demand, train, month_weekdays, q_SAA=None):
    """
    Create a stripplot of the demand per month and day of the week for the first two months.

    Parameters:
    features (DataFrame): DataFrame containing feature data.
    demand (Series): Series containing demand data.
    train (array-like): Indices or boolean mask for training data.
    month_weekdays (list): List of month weekdays in the desired order.
    q_SAA (float, Series, or None): SAA quantity. If None, no line is added.
    Returns:
    None
    """
    # Set up the figure and axis
    plt.figure(figsize=(15, 8))
    
    # Create the stripplot
    ax = sns.stripplot(x=features['month_weekday'][train], y=demand[train], jitter=True, dodge=True, order=month_weekdays[:14])

    # Add a horizontal line for the SAA quantity if it is provided
    if q_SAA is not None:

        # Check if this a single value or an array / pandas Series
        if isinstance(q_SAA, (int, float)):
            plt.axhline(y=q_SAA, color='orange', linestyle='-', label="SAA", alpha=0.5)
        else:
            for i, month_weekday in enumerate(month_weekdays[:14]):
                q_month_day = q_SAA.loc[month_weekday]
                plt.axhline(y=q_month_day, color='orange', linestyle='-', alpha=0.5, xmin=i/14, xmax=(i+1)/14)

            plt.plot([], [], color='orange', linestyle='-', alpha=0.5, label="SAA per day")
    
        # Add the legend
        plt.legend()
    
    # Set the title
    plt.title('Demand by day/month (only first two months shown)', fontsize=16)
    
    # Display the plot
    plt.show()





def plot_prescriptions_weekday_order_quantities(order_quantities, features_train, demand_train):
    colors = ['tab:blue', 'orange', 'green']
    line_styles = ['-', '--', 'dotted']

    plt.figure(figsize=(10,6))
    ax = sns.stripplot(x=features_train['weekday'],
                        y=demand_train,
                        jitter=True, dodge=True, order = weekdays, alpha=0.5)
    
    for i, day in enumerate(weekdays):
        for j, method in enumerate(order_quantities.index):
    
    
            # make thick line for the optimal quantity
            plt.axhline(
                y=order_quantities.T[method].loc[day],
                color=colors[j],
                linestyle=line_styles[j],
                linewidth=2.5,
                xmin=i/7,
                xmax=(i+1)/7,
                label=method if i == 0 else None
            )
    plt.legend()
    plt.show()
    return 




def plot_prescriptions_weekday(lr, lrnw,  weekday_train, features_train, demand_train):


    # get the unique weekdays
    weekdays_unique = weekday_train[3:][:7].set_axis(weekdays)

    q_lr = lr.predict(weekdays_unique)
    q_lrnw = lrnw.predict(weekdays_unique)

    order_quantities = pd.DataFrame({'Quantity': q_lr.flatten(), 'Quantity Newsvendor': q_lrnw.flatten()}, index=weekdays).T
    
    plot_prescriptions_weekday_order_quantities(order_quantities, features_train, demand_train)
    


service_levels = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
service_level = cu / (cu + co)




def interactive_weekday_plot(lr,  weekday_train, features_train, demand_train):


    # Ignore all warnings in python
    import warnings

    warnings.filterwarnings('ignore')
    



    # get the unique weekdays
    weekdays_unique = weekday_train[3:][:7].set_axis(weekdays)
    q_lr = lr.predict(weekdays_unique)

    # Create the order_quantity_service_level dataframe
    order_quantity_service_level = pd.DataFrame(columns=weekdays)
    # Add the linear regression results to the order_quantity_service_level dataframe using pandas concat. (for each weekday)
    order_quantity_service_level = pd.concat([order_quantity_service_level,
                                                pd.DataFrame(q_lr.reshape(1, -1), columns=weekdays, index=["Linear regression"])])
    

    for service_level in service_levels:

        # create the quantile regression model and fit it to the data
        qr = QuantileRegressor(quantile=service_level, alpha=0, solver='highs').fit(weekday_train, demand_train)
        # predict the optimal quantities for the unique weekdays
        q_qr_unique = np.round(qr.predict(weekdays_unique))
        # Add the optimal quantities to the order_quantity_service_level dataframe using pandas concat
        order_quantity_service_level = pd.concat([
                                                    pd.DataFrame(q_qr_unique.reshape(1, -1), columns=weekdays, index=[f"Linear regression newsvendor service level={service_level}"]),
                                                    order_quantity_service_level
                                                    ])
        


    
    service_level_slider = widgets.FloatSlider(
    value=service_levels[0],
    min=service_levels[0],
    max=service_levels[-1],
    step=0.1,
    description='Service level:',
    disabled=False,
    continuous_update=False,
    orientation='horizontal',
    readout=True,
    readout_format='.2f',
    )

    # For the selected service level, plot the
    # optimal quantities of the linear regression newsvendor (f"Linear regression newsvendor service level={sl}") and the linear regression model

    def plot_optimal_quantities(service_level):
        # select the follwing rows from the order_quantity_service_level dataframe: f"Linear regression newsvendor service level={sl}" and "Linear regression"
        service_level = round(service_level, 2)
        indices = [ "Linear regression", f"Linear regression newsvendor service level={service_level}" ]
        df = order_quantity_service_level.loc[indices]
        plot_prescriptions_weekday_order_quantities(df, features_train, demand_train)
        return
    

    return widgets.interact(plot_optimal_quantities, service_level=service_level_slider)





def plot_lrnw_core(data, demand, feature, saa_predictions, lrnw_predictions, lr_predictions):

    # create a new plot using matplotlib
    plt.figure(figsize=(12, 6))

    # Plot the demand against the temperature show demand 0 to 50
    plt.scatter(data[feature], demand, alpha=0.5)
    # Add axis labels
    plt.xlabel(feature)
    plt.ylabel('Demand')


    if saa_predictions is not None:
        # Add horizontal line using the linear regression model
        x = np.linspace(data[feature].min(), data[feature].max(), len(demand))
        plt.plot(x, saa_predictions, color='darkred', label='SAA', linestyle='--', linewidth=2)
    # Add regression line using the linear regression newsvendor model

    if lrnw_predictions is not None:
        plt.plot(data[feature], lrnw_predictions, color='orange', label='Linear regression newsvendor', linewidth=1.75)

    if lr_predictions is not None:
        plt.plot(data[feature], lr_predictions, color='navy', label='Linear regression', linewidth=1.75)

    # Make the minimum of the y-axis -5 and the maximum 100
    plt.ylim(-5, 100)


    # Add headline to the plot
    plt.title(f"Demand against {feature}")
    # Add legend
    plt.legend()

    return


def plot_lin_reg(data, demand, feature, models):

    if 'SAA' in models:
        saa_quantity = 35
        saa_predictions = np.full_like(demand, saa_quantity)
    else :
        saa_predictions = None

    if 'LRNW' in models:
        lrnw.fit(data[feature].to_frame(), demand)
        lrnw_predictions = lrnw.predict(data[feature].to_frame()).flatten()
    else:
        lrnw_predictions = None

    if 'LR' in models:
        lr.fit(data[feature].to_frame(), demand)
        lr_predictions = lr.predict(data[feature].to_frame()).flatten()
    else:
        lr_predictions = None

    plot_lrnw_core(data, demand, feature, saa_predictions, lrnw_predictions, lr_predictions)
    return


def plot_loss_functions():

    # Plot the loss function for the linear regression model (square loss)
    # Plot the loss function for the linear regression newsvendor model (co / cu loss)

    # Plot a quadratic loss function around 0

    x = np.linspace(-10, 10, 100)
    y_squared = x**2
    y_newsvendor = np.where(x < 0, cu * -x, co * x)

    plt.figure(figsize=(12, 6))
    sns.lineplot(x=x, y=y_squared, label="Square Loss")
    sns.lineplot(x=x, y=y_newsvendor, label="Newsvendor Loss (cu, co)")
    plt.xlabel("Prediction Error")
    plt.ylabel("Loss")
    plt.title("Loss Functions")
