#!/usr/bin/env python
# coding: utf-8

##################################################### IMPORT PACKAGES #################################################

import os
import csv
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.image as img
import matplotlib.ticker as tick
import seaborn as sns
import ipywidgets as wg
from ipywidgets import interact, interact_manual, Button, Layout, Label, HBox, VBox
from IPython.display import display

import collections


##################################################### CASE CLASS #################################################


class CaseHandler:
    # Dictionary to record all instances of Case class (key = case directory (case_name), item = Case instance object)
    active_cases = {}

    # Important variables common to all CaseHandler class instances
    # Relative path to inputs and results directories
    # can manually change with CaseHandler.inputs_directory = <inputs directory path> then calling CaseHandler.update()
    inputs_directory = os.path.join("inputs")
    results_directory = os.path.join("results")

    # Important widget output object
    output = wg.Output()

    @classmethod
    def update(cls):

        """
        Updates variables common to all CaseHandler class instances with respect to results_directory and inputs_directory
        """
        # Dataframe of color_map.csv file for mapping colors to resource groups
        cls.color_map = pd.read_csv(os.path.join("color_map.csv"))

    @classmethod
    def load_cases(cls):

        """
        Call "CaseHandler.select_cases()"" to create a checkbox list of cases in results_directory
        Select the case check boxes and click the "Load Selected Cases" button to create CaseHandler instances for selected cases

        Currently, there is no simple way to remove a CaseHandler instance from "active_cases"
        Once a Case instance is created/loaded, it will remain in the active_cases dictionary
        """
        print("Select cases to load:")
        # Update variables common to all class instances (based on results_directory and inputs_directory)
        cls.update()
        # Get list of all cases in results_directory
        cls.all_cases = os.listdir(cls.results_directory)
        cls.case_check_boxes = [
            wg.Checkbox(value=False, description=case_name, disabled=False)
            for case_name in cls.all_cases
        ]
        cls.create_cases_button = wg.Button(
            description="Load Selected Cases",
            layout=Layout(width="90%"),
            disabled=False,
        )
        cls.create_cases_button.on_click(cls.create_cases)
        display(wg.VBox(cls.case_check_boxes))
        display(cls.create_cases_button)
        display(cls.output)

    @classmethod
    def create_cases(cls, args=None):

        """
        Iteratively creates CaseHandler instances for each new case name in CaseHandler.selected_case_names
        Gets called upon clicking the "Load Selected Cases" button (create_cases_button)
        """

        cls.selected_case_names = [
            c.description for c in cls.case_check_boxes if c.value
        ]
        with cls.output:
            for case_name in cls.selected_case_names:
                # Do not re-load previously selected cases
                if case_name not in cls.active_cases.keys():
                    print("Loading new case...")
                    # Create Case instance from case_name
                    cls(case_name)

    def __init__(self, case_name):

        # Add new instance to active_cases; case_name should be name of case directory found in results_directory
        CaseHandler.active_cases[case_name] = self
        print("Loaded case: {}\n".format(case_name))

        # Initialize variables associated with case

        self.case_name = case_name
        self.inputs_folder = os.path.join(
            CaseHandler.inputs_directory, case_name
        )  # Path to inputs folder for case
        self.results_folder = os.path.join(
            CaseHandler.results_directory, case_name
        )  # Path to results folder for case

        # Dataframe of the resource_map.csv file which is used to construct type_map and type_order
        self.resource_map = pd.read_csv(
            os.path.join(self.results_folder, "resource_results_grouping.csv")
        )
        # List of zones in RESOLVE model
        self.zones = self.resource_map.zone.unique()
        # Resources grouped by zone... MAY BE OUTDATED NOW AND UNNECESSARY
        self.zone_resources = self.resource_map.groupby(["zone", "results_grouping"])
        # type_map is dictionary with key (zone, resource group) and gives list of resources for that group
        # ... MAY BE OUTDATED NOW AND UNNECESSARY
        self.type_map = {
            group: self.zone_resources.get_group(group)["resource"].values.tolist()
            for group in self.zone_resources.groups
        }
        # type_order is dictionary with key (zone, resource group) and gives the plotting order for that group
        # ... MAY BE OUTDATED NOW AND UNNECESSARY
        type_order = {
            group: self.zone_resources.get_group(group)["group_order"].iloc[0]
            for group in self.zone_resources.groups
        }

        # Create accessible dataframes for key output files
        self.operations = pd.read_csv(
            os.path.join(self.results_folder, "operations.csv")
        ).fillna(0)
        self.operational_costs = pd.read_csv(
            os.path.join(self.results_folder, "summary", "operational_costs.csv")
        ).fillna(0)
        self.loads_and_power_balance = pd.read_csv(
            os.path.join(self.results_folder, "loads_and_power_balance.csv")
        ).fillna(0)
        self.load = pd.read_csv(
            os.path.join(self.results_folder, "summary", "annual_load.csv")
        ).fillna(0)
        self.transmit_power = pd.read_csv(
            os.path.join(self.results_folder, "transmit_power.csv")
        ).fillna(0)
        self.transmission_flows = pd.read_csv(
            os.path.join(self.results_folder, "summary", "transmission_flows.csv")
        ).fillna(0)
        self.resource_build = pd.read_csv(
            os.path.join(self.results_folder, "resource_build.csv")
        ).fillna(0)
        self.resource_energy = pd.read_csv(
            os.path.join(self.results_folder, "summary", "annual_energy.csv")
        ).fillna(0)
        self.emissions = pd.read_csv(
            os.path.join(self.results_folder, "summary", "emissions.csv")
        ).fillna(0)
        self.reserve_timepoints = pd.read_csv(
            os.path.join(self.results_folder, "reserve_timepoints.csv")
        ).fillna(0)
        self.reliability_constraints = pd.read_csv(
            os.path.join(self.results_folder, "reliability_constraints.csv")
        ).fillna(0)

        # Create accessible dataframes for key input files
        # self.resource_results_grouping = pd.read_csv(os.path.join(self.inputs_folder, 'resource_results_grouping.csv'))
        self.timepoints = self.reserve_timepoints["timepoint"].unique()

        # Define important variables from accessible dataframes
        # Representative days in RESOLVE model
        self.days = self.reserve_timepoints["day"].unique()
        # List of years in RESOLVE model
        self.years = self.reserve_timepoints["period"].unique()
        # List of resource types in RESOLVE model
        self.resource_types = self.resource_map.sort_values(
            by=["group_order"]
        ).results_grouping.unique()

        # Run some data pre-processing steps
        self.get_resource_types()
        self.group_resources_by_type()

    def get_resource_types(self):
        # Add column to self.resource_build dataframe that includes resource types defined by self.resource_map

        map_df = self.resource_map.filter(["resource", "results_grouping"]).set_index(
            "resource"
        )["results_grouping"]
        self.resource_build["resource_type"] = self.resource_build["resource"].map(
            map_df
        )
        self.resource_energy["resource_type"] = self.resource_energy["resource"].map(
            map_df
        )
        self.operations["resource_type"] = self.operations["resource"].map(map_df)
        self.emissions["resource_type"] = self.emissions["resource"].map(map_df)
        self.operational_costs["resource_type"] = self.operational_costs[
            "resource"
        ].map(map_df)

    def group_resources_by_type(self):

        self.resources = (
            pd.pivot_table(
                self.resource_build,
                values=[
                    "operational_capacity_mw",
                    "cumulative_new_build_mw",
                    "retired_planned_capacity_mw",
                ],
                index=["period", "resource_type", "zone"],
                aggfunc="sum",
            )
            .sort_values(by=["period", "zone", "resource_type"])
            .reset_index()
        )

        self.resources[
            ["net_energy_mwh_per_year", "scheduled_curtailment_mwh_per_year"]
        ] = (
            pd.pivot_table(
                self.resource_energy,
                values=[
                    "net_energy_mwh_per_year",
                    "scheduled_curtailment_mwh_per_year",
                ],
                index=["period", "resource_type", "zone"],
                aggfunc="sum",
            )
            .sort_values(by=["period", "zone", "resource_type"])
            .reset_index()[
                ["net_energy_mwh_per_year", "scheduled_curtailment_mwh_per_year"]
            ]
        )

        self.resources[list(self.emissions.emission_type.unique())] = (
            pd.pivot_table(
                self.emissions,
                values=["emissions_units_per_year"],
                index=["period", "resource_type", "zone"],
                columns="emission_type",
                aggfunc="sum",
            )
            .sort_values(by=["period", "zone", "resource_type"])
            .reset_index()["emissions_units_per_year"][
                list(self.emissions.emission_type.unique())
            ]
        )

        self.resources[
            [
                "total_fr_mw_per_year",
                "spin_mw_per_year",
                "upward_reg_mw_per_year",
                "downward_reg_mw_per_year",
                "upward_lf_mw_per_year",
                "downward_lf_mw_per_year",
            ]
        ] = (
            pd.pivot_table(
                self.operational_costs,
                values=[
                    "total_fr_mw_per_year",
                    "spin_mw_per_year",
                    "upward_reg_mw_per_year",
                    "downward_reg_mw_per_year",
                    "upward_lf_mw_per_year",
                    "downward_lf_mw_per_year",
                ],
                index=["period", "resource_type", "zone"],
                aggfunc="sum",
            )
            .sort_values(by=["period", "zone", "resource_type"])
            .reset_index()[
                [
                    "total_fr_mw_per_year",
                    "spin_mw_per_year",
                    "upward_reg_mw_per_year",
                    "downward_reg_mw_per_year",
                    "upward_lf_mw_per_year",
                    "downward_lf_mw_per_year",
                ]
            ]
        )

        self.dispatch = (
            pd.pivot_table(
                self.operations,
                values=[
                    "power_mw",
                    "scheduled_curtailment_mw",
                    "increase_load_mw",
                    "net_power_mw",
                ],
                index=["period", "day", "hour_of_day", "resource_type", "zone"],
                aggfunc="sum",
            )
            .sort_values(by=["period", "day", "hour_of_day", "zone", "resource_type"])
            .reset_index()
        )


####################################### PLOT CLASS ###################################################


class Plot:

    # Initialization and update functions

    def __init__(self):

        self.style = {"description_width": "initial"}  # Style for widgets

        self.caselist = wg.Dropdown(
            options=CaseHandler.active_cases.keys(),
            description="Case:",
            style=self.style,
            disabled=False,
        )  # Dropdown selection of cases
        self.caselist.observe(self.update, names="value")  # Observe and update
        self.case = CaseHandler.active_cases[self.caselist.value]

        self.caseslist = wg.SelectMultiple(
            options=CaseHandler.active_cases.keys(),
            description="Cases:",
            style=self.style,
            disabled=False,
        )  # Multiple selection of cases
        self.caseslist.observe(self.update, names="value")  # Observe and update
        self.cases = [CaseHandler.active_cases[val] for val in self.caseslist.value]

        self.plot = None  # Function to plot on clicking of plot button
        self.plot_button = Button(
            description="Plot", layout=Layout(width="90%"), disabled=False
        )  # Plot button

        self.years = wg.SelectMultiple(
            options=self.case.years,
            description="Years:",
            style=self.style,
            disabled=False,
        )
        self.year = wg.Dropdown(
            options=self.case.years,
            description="Year:",
            style=self.style,
            disabled=False,
        )
        self.zones = wg.SelectMultiple(
            options=self.case.zones,
            description="Zones:",
            style=self.style,
            disabled=False,
        )
        self.days = wg.SelectMultiple(
            options=self.case.days,
            description="Days:",
            style=self.style,
            disabled=False,
        )
        self.col = wg.Dropdown(
            options=[
                "operational_capacity_mw",
                "cumulative_new_build_mw",
                "retired_planned_capacity_mw",
                "net_energy_mwh_per_year",
                "scheduled_curtailment_mwh_per_year",
                "total_fr_mw_per_year",
                "spin_mw_per_year",
                "upward_reg_mw_per_year",
                "downward_reg_mw_per_year",
                "upward_lf_mw_per_year",
                "downward_lf_mw_per_year",
            ]
            + list(self.case.emissions.emission_type.unique()),
            description="Column:",
            style=self.style,
            disabled=False,
        )

        self.hmap_col = wg.Dropdown(
            options=[
                "hourly_energy_cost_$_per_mwh",
                "input_load_mw",
                "net_power_mw",
                "scheduled_curtailment_mw",
                "unspecified_power_mw",
            ],
            description="Column:",
            style=self.style,
            disabled=False,
        )

        self.resource_types = wg.SelectMultiple(
            options=self.case.resource_types,
            description="Resources Types (if applicable):",
            style=self.style,
            disabled=False,
        )

        self.tx_line = wg.Dropdown(
            options=self.case.transmit_power.transmission_line.unique(),
            description="TX Line (if applicable):",
            style=self.style,
            disabled=False,
        )

        self.output = wg.Output()

    def update(self, args=None):
        with self.output:
            self.output.clear_output()
            self.case = CaseHandler.active_cases[self.caselist.value]
            self.cases = [CaseHandler.active_cases[val] for val in self.caseslist.value]
            self.years.options = self.case.years
            self.year.options = self.case.years
            self.zones.options = self.case.zones
            self.days.options = self.case.days
            self.col.options = [
                "operational_capacity_mw",
                "cumulative_new_build_mw",
                "retired_planned_capacity_mw",
                "net_energy_mwh_per_year",
                "scheduled_curtailment_mwh_per_year",
                "total_fr_mw_per_year",
                "spin_mw_per_year",
                "upward_reg_mw_per_year",
                "downward_reg_mw_per_year",
                "upward_lf_mw_per_year",
                "downward_lf_mw_per_year",
            ] + list(self.case.emissions.emission_type.unique())
            self.hmap_col.options = [
                "hourly_energy_cost_$_per_mwh",
                "input_load_mw",
                "net_power_mw",
                "scheduled_curtailment_mw",
                "unspecified_power_mw",
            ]
            self.resource_types.options = self.case.resource_types
            self.tx_line.options = self.case.transmit_power.transmission_line.unique()
            self.plot_button.on_click(self.plot)

    # Widget functions

    def select_case(self):
        display(self.caselist)

    def select_multiple_cases(self):
        display(self.caseslist)

    def select_years(self):
        display(self.years)

    def select_year(self):
        display(self.year)

    def select_zones(self):
        display(self.zones)

    def select_days(self):
        display(self.days)

    def select_column(self):
        display(self.col)

    def select_hmap_column(self):
        display(self.hmap_col)

    def select_resource_types(self):
        display(self.resource_types)

    def select_tx_line(self):
        display(self.tx_line)

    def get_plot_button(self):
        self.plot_button.on_click(self.plot)
        display(self.plot_button)
        display(self.output)

    # Plotting functions

    def annual_resource_plot(self, args=None):
        with self.output:
            self.output.clear_output()

            ####### INITIALIZE PLOT AND DATA #######
            # initialize the figure, x axis (years), and bar stack array
            x = np.arange(len(self.years.value) * len(self.caseslist.value))
            yr = np.arange(len(self.years.value))
            width = 0.75
            fig, ax = plt.subplots(1, figsize=(8, 6))
            bar_arr = []

            ####### PLOT RESOURCE BARS #######
            c = 0
            for case in self.caseslist.value:

                # initialize the bottom trackers for stacking
                bottom = np.zeros(
                    len(self.years.value)
                )  # Initialize bottom of stacked bar chart
                neg_bottom = np.zeros(
                    len(self.years.value)
                )  # Initialize negative bottom of stacked bar chart (if applicable)
                pos_values = np.zeros(len(self.years.value))  # Initialize value array
                neg_values = np.zeros(len(self.years.value))  # Initialize value array

                # gather the resource data for the annual plot based on user inputs for years and zone aggregation
                filtered_data = CaseHandler.active_cases[case].resources.loc[
                    (
                        CaseHandler.active_cases[case]
                        .resources["period"]
                        .isin(self.years.value)
                    )
                    & (
                        CaseHandler.active_cases[case]
                        .resources["zone"]
                        .isin(self.zones.value)
                    )
                ]
                pivot_data = pd.pivot_table(
                    filtered_data,
                    values=self.col.value,
                    index=["period"],
                    columns=["resource_type"],
                    aggfunc="sum",
                )

                # for each resource type in order, plot a bar and if the resource type doesn't exist in the data, pass over
                for resource_name in CaseHandler.active_cases[case].resource_types:
                    if resource_name not in pivot_data.columns:
                        continue

                    # in case of negative annual net energy (i.e. storage), need to track and add to figure separately
                    pos_values = pivot_data[resource_name].values.clip(min=0)
                    neg_values = pivot_data[resource_name].values.clip(max=0)
                    if c == 0:
                        bar_arr.append(
                            ax.bar(
                                yr + c * len(self.years.value),
                                pos_values,
                                width,
                                color=CaseHandler.active_cases[case].color_map.loc[
                                    CaseHandler.active_cases[case].color_map[
                                        "resource type"
                                    ]
                                    == resource_name,
                                    "color",
                                ],
                                edgecolor="w",
                                bottom=bottom,
                                label=resource_name,
                            )
                        )
                        ax.bar(
                            yr + c * len(self.years.value),
                            neg_values,
                            width,
                            color=CaseHandler.active_cases[case].color_map.loc[
                                CaseHandler.active_cases[case].color_map[
                                    "resource type"
                                ]
                                == resource_name,
                                "color",
                            ],
                            edgecolor="w",
                            bottom=neg_bottom,
                            label=resource_name,
                        )
                    else:
                        ax.bar(
                            yr + c * len(self.years.value),
                            pos_values,
                            width,
                            color=CaseHandler.active_cases[case].color_map.loc[
                                CaseHandler.active_cases[case].color_map[
                                    "resource type"
                                ]
                                == resource_name,
                                "color",
                            ],
                            edgecolor="w",
                            bottom=bottom,
                            label=resource_name,
                        )
                        ax.bar(
                            yr + c * len(self.years.value),
                            neg_values,
                            width,
                            color=CaseHandler.active_cases[case].color_map.loc[
                                CaseHandler.active_cases[case].color_map[
                                    "resource type"
                                ]
                                == resource_name,
                                "color",
                            ],
                            edgecolor="w",
                            bottom=neg_bottom,
                            label=resource_name,
                        )

                    # update bottoms of the stack
                    bottom = bottom + pos_values
                    neg_bottom = neg_bottom + neg_values

                # in case of plotting annual energy, want to also plot imports and curtailment
                if self.col.value == "net_energy_mwh_per_year":

                    ###### PLOT IMPORTS ######
                    # filter the transmission data by user-defined period
                    filtered_tx = self.case.transmission_flows.loc[
                        (
                            CaseHandler.active_cases[case]
                            .transmission_flows["period"]
                            .isin(self.years.value)
                        )
                    ]

                    # track both directions for net flows (from and to direction)
                    tx_from = filtered_tx[
                        filtered_tx.transmission_from.isin(self.zones.value)
                    ]
                    tx_to = filtered_tx[
                        filtered_tx.transmission_to.isin(self.zones.value)
                    ]

                    # if zone(s) doesn't exist for one direction, make an array of zeros for the that direction
                    # otherwise, pivot the power (net so can be +/-) by adding up flows for zone(s) of interest
                    if tx_from.empty:
                        from_imports = np.zeros_like(self.years.value)
                    else:
                        from_imports = pd.pivot_table(
                            tx_from,
                            values="power_mwh_per_year",
                            index=["period"],
                            aggfunc="sum",
                        )["power_mwh_per_year"].values

                    if tx_to.empty:
                        to_imports = np.zeros_like(self.years.value)
                    else:
                        to_imports = pd.pivot_table(
                            tx_to,
                            values="power_mwh_per_year",
                            index=["period"],
                            aggfunc="sum",
                        )["power_mwh_per_year"].values

                    # imports take difference of aggregate flows for to and from directions - only positive values imports
                    imports = (to_imports - from_imports).clip(min=0)

                    # plot imports as bar and update bottom stack
                    if c == 0:

                        bar_arr.append(
                            ax.bar(
                                yr + c * len(self.years.value),
                                imports,
                                width,
                                color=CaseHandler.active_cases[case].color_map.loc[
                                    CaseHandler.active_cases[case].color_map[
                                        "resource type"
                                    ]
                                    == "Imports",
                                    "color",
                                ],
                                edgecolor="w",
                                bottom=bottom,
                                label="Imports",
                            )
                        )
                    else:
                        ax.bar(
                            yr + c * len(self.years.value),
                            imports,
                            width,
                            color=CaseHandler.active_cases[case].color_map.loc[
                                CaseHandler.active_cases[case].color_map[
                                    "resource type"
                                ]
                                == "Imports",
                                "color",
                            ],
                            edgecolor="w",
                            bottom=bottom,
                            label="Imports",
                        )
                    bottom = bottom + imports

                    ######## PLOT CURTAILMENT ########
                    # get curtailment data and add as bar to plot, update bottom stack
                    curtailment = pd.pivot_table(
                        filtered_data,
                        values="scheduled_curtailment_mwh_per_year",
                        index=["period"],
                        aggfunc="sum",
                    )["scheduled_curtailment_mwh_per_year"].values

                    if c == 0:
                        bar_arr.append(
                            ax.bar(
                                yr + c * len(self.years.value),
                                curtailment,
                                width,
                                color=CaseHandler.active_cases[case].color_map.loc[
                                    CaseHandler.active_cases[case].color_map[
                                        "resource type"
                                    ]
                                    == "Curtailment",
                                    "color",
                                ],
                                edgecolor="w",
                                bottom=bottom,
                                label="Curtailment",
                            )
                        )
                    else:
                        ax.bar(
                            yr + c * len(self.years.value),
                            curtailment,
                            width,
                            color=CaseHandler.active_cases[case].color_map.loc[
                                CaseHandler.active_cases[case].color_map[
                                    "resource type"
                                ]
                                == "Curtailment",
                                "color",
                            ],
                            edgecolor="w",
                            bottom=bottom,
                            label="Curtailment",
                        )

                    bottom = bottom + curtailment

                    ######## PLOT LOAD MARKER ########
                    # get load data and add as marker on plot
                    load_line = pd.pivot_table(
                        CaseHandler.active_cases[case].load.loc[
                            (
                                CaseHandler.active_cases[case]
                                .load["period"]
                                .isin(self.years.value)
                            )
                            & (
                                CaseHandler.active_cases[case]
                                .load["zone"]
                                .isin(self.zones.value)
                            )
                        ],
                        values="input_load_mwh_per_year",
                        index="period",
                        aggfunc="sum",
                    )["input_load_mwh_per_year"].values
                    (line,) = ax.plot(
                        yr + c * len(self.years.value),
                        load_line,
                        marker="D",
                        markerfacecolor="black",
                        markeredgecolor="black",
                        linestyle="None",
                        label="Load",
                    )

                c += 1

            ######## SET FIGURE OPTIONS ########

            ax.set_xticks(x, minor=True)

            ax.set_xticklabels(self.years.value * len(self.caseslist.value), minor=True)

            ax.set_xticks(
                np.arange(len(self.caseslist.value)) * len(self.years.value)
                + (len(self.years.value) - 1) / 2
                + 0.001
            )
            ax.set_xticklabels(self.caseslist.value)
            ax.tick_params(axis="x", pad=20, length=0)

            if len(self.caseslist.value) > 0:
                for c in np.arange((len(self.caseslist.value) - 1)):
                    ax.axvline(
                        x=(c + 1) * (len(self.years.value)) - 0.5, color="gray"
                    )

            ax.set_ylim([min(neg_bottom) * 1.1, max(bottom) * 1.1])

            # if the figure shows capacity, set as y axis label - otherwise it would be energy so set that y axis label
            if self.col.value in [
                "operational_capacity_mw",
                "cumulative_new_build_mw",
                "retired_planned_capacity_mw",
            ]:
                ax.set_ylabel("Capacity (MW)")
            elif self.col.value in list(
                CaseHandler.active_cases[
                    self.caseslist.value[0]
                ].emissions.emission_type.unique()
            ):
                ax.set_ylabel("Emissions (units)")
            else:
                ax.set_ylabel("Energy (MWh)")

            ax.set_title("Portfolio by Year")

            # if annual energy plot is selected, need to add the load marker line to the legend
            if self.col.value == "net_energy_mwh_per_year":
                ax.legend(
                    handles=bar_arr[::-1] + [line],
                    bbox_to_anchor=[1, 1],
                    loc="upper left",
                    labelspacing=1.05,
                    frameon=False,
                )
            else:
                ax.legend(
                    handles=bar_arr[::-1],
                    bbox_to_anchor=[1, 1],
                    loc="upper left",
                    labelspacing=1.05,
                    frameon=False,
                )

            fig.tight_layout()

            plt.show()

    def dispatch_plot(self, args=None):
        with self.output:
            self.output.clear_output()

            ####### INITIALIZE PLOT AND DATA #######
            # initialize the figure as a grid of suplots, x axis (days) and y axis (years)
            fig, ax = plt.subplots(
                len(self.years.value), len(self.days.value), sharex=True, sharey=True
            )
            fig.set_tight_layout(False)
            fig.subplots_adjust(
                left=0.125, bottom=0.1, right=0.9, top=0.9, wspace=0.1, hspace=0.2
            )
            leg_x = 1.1
            leg_y = 1.1

            ###### PLOT DATA ON EACH SUBPLOT #######
            # loop through each selected day and year and add to the subplot
            for j, representative_day in enumerate(self.days.value):
                for i, year in enumerate(self.years.value):

                    # defining different sets of axes for each subplot depending on selected days and years
                    if len(self.days.value) > 1 and len(self.years.value) > 1:
                        ax_to_plot = ax[i, j]
                        ax_stor = ax[-1, 0]
                        leg_x = len(self.days.value) + 0.1 * len(self.days.value)
                        leg_y = len(self.years.value) + 0.125 * len(self.years.value)
                    elif len(self.years.value) > 1:
                        ax_to_plot = ax[i]
                        ax_stor = ax[0]
                    elif len(self.days.value) > 1:
                        ax_to_plot = ax[j]
                        ax_stor = ax[-1]
                    else:
                        ax_to_plot = ax
                        ax_stor = ax

                    # define the legend handles and the bottom tracker
                    legend_handles = []
                    bottom = np.zeros(24)

                    # gather the resource generation data based on user inputs for year, day, and zone aggregation
                    filtered_data = self.case.dispatch.loc[
                        (self.case.dispatch["period"] == year)
                        & (self.case.dispatch["zone"].isin(self.zones.value))
                        & (self.case.dispatch["day"] == representative_day)
                    ]

                    # get the zonal aggregate values with columns as resource types
                    pivot_data = pd.pivot_table(
                        filtered_data,
                        values="power_mw",
                        index=["period", "day", "hour_of_day"],
                        columns=["resource_type"],
                        aggfunc="sum",
                    ).reset_index()

                    ####### PLOT RESOURCE GENERATION #######
                    # for each resource type in order, plot a stacked area if the resource type exists in the data
                    for resource_name in self.case.resource_types:
                        if resource_name not in pivot_data.columns:
                            continue

                        # add values to the subplot and update the bottom of the stack
                        values = pivot_data[resource_name].values
                        legend_handles.append(
                            ax_to_plot.fill_between(
                                np.arange(24),
                                bottom,
                                bottom + values,
                                linewidth=0,
                                color=self.case.color_map.loc[
                                    self.case.color_map["resource type"]
                                    == resource_name,
                                    "color",
                                ],
                                label=resource_name,
                            )
                        )
                        ax_to_plot.plot(np.arange(24), bottom, linewidth=0, color="w")
                        bottom += values

                    ######## PLOT IMPORTS ########
                    # filter the transmission data by user-defined year and day
                    filtered_tx = self.case.transmit_power.loc[
                        (self.case.transmit_power["period"] == year)
                        & (self.case.transmit_power["day"] == representative_day)
                    ]

                    # track both directions for net flows (from and to direction)
                    tx_from = filtered_tx[
                        filtered_tx.transmission_from.isin(self.zones.value)
                    ]
                    tx_to = filtered_tx[
                        filtered_tx.transmission_to.isin(self.zones.value)
                    ]

                    # if zone(s) doesn't exist for one direction, make an array of zeros for the that direction
                    # otherwise, pivot the power (net so can be +/-) by adding up flows for zone(s) of interest
                    if tx_from.empty:
                        from_imports = np.zeros(24)
                    else:
                        from_imports = pd.pivot_table(
                            tx_from,
                            values="unspecified_power_mw",
                            index=["period", "day", "hour_of_day"],
                            aggfunc="sum",
                        )["unspecified_power_mw"].values

                    if tx_to.empty:
                        to_imports = np.zeros(24)
                    else:
                        to_imports = pd.pivot_table(
                            tx_to,
                            values="unspecified_power_mw",
                            index=["period", "day", "hour_of_day"],
                            aggfunc="sum",
                        )["unspecified_power_mw"].values

                    # imports take difference of aggregate flows in to and from directions - positive values are imports
                    imports = (to_imports - from_imports).clip(min=0)
                    exports = (from_imports - to_imports).clip(min=0)

                    # plot the import energy
                    legend_handles.append(
                        ax_to_plot.fill_between(
                            np.arange(24),
                            bottom,
                            bottom + imports,
                            linewidth=0,
                            facecolor=self.case.color_map.loc[
                                self.case.color_map["resource type"] == "Imports",
                                "color",
                            ],
                            label="Imports",
                        )
                    )

                    ax_to_plot.plot(np.arange(24), bottom, linewidth=0, color="w")
                    bottom += imports

                    ######## PLOT CURTAILMENT ########
                    # get aggregate curtailment for year, day, and zone(s) and add to subplot, update bottom stack
                    curtailment_pivot = pd.pivot_table(
                        filtered_data,
                        values="scheduled_curtailment_mw",
                        index=["period", "day", "hour_of_day"],
                        aggfunc="sum",
                    ).reset_index()
                    curtailment = curtailment_pivot["scheduled_curtailment_mw"].values
                    legend_handles.append(
                        ax_to_plot.fill_between(
                            np.arange(24),
                            bottom,
                            bottom + curtailment,
                            linewidth=0,
                            facecolor=self.case.color_map.loc[
                                self.case.color_map["resource type"] == "Curtailment",
                                "color",
                            ],
                            label="Curtailment",
                        )
                    )
                    ax_to_plot.plot(np.arange(24), bottom, linewidth=0, color="w")
                    bottom += curtailment

                    ######## PLOT LOAD LINES ########
                    # get filtered load data for specific day, year, and zone(s)
                    filtered_load = self.case.loads_and_power_balance.loc[
                        (self.case.loads_and_power_balance["period"] == year)
                        & (
                            self.case.loads_and_power_balance["zone"].isin(
                                self.zones.value
                            )
                        )
                        & (
                            self.case.loads_and_power_balance["day_id"]
                            == representative_day
                        )
                    ]

                    # pivot to sum up load across different zones
                    load_pivot = pd.pivot_table(
                        filtered_load,
                        values="input_load_mw",
                        index=["period", "day_id", "hour_of_day"],
                        aggfunc="sum",
                    ).reset_index()

                    # add load as solid line to the plot
                    load = load_pivot["input_load_mw"].values
                    (ln_load,) = ax_to_plot.plot(
                        np.arange(24), load, label="Load", color="k", linewidth=1.5
                    )

                    # pivot to sum up storage charging across different zones
                    storage_charging = pd.pivot_table(
                        filtered_load,
                        values="increase_load_mw",
                        index=["period", "day_id", "hour_of_day"],
                        aggfunc="sum",
                    ).reset_index()

                    # add load + storage charging as dashed line
                    load += storage_charging["increase_load_mw"].values
                    (ln_loadst,) = ax_to_plot.plot(
                        np.arange(24),
                        load,
                        color="k",
                        linewidth=1.5,
                        label="Load + Storage charge",
                        dashes=[5, 5],
                    )

                    load += exports
                    (ln_loadexp,) = ax_to_plot.plot(
                        np.arange(24),
                        load,
                        color="k",
                        linewidth=2,
                        label="Load + Storage charge + exports",
                        dashes=[1.5, 1.5],
                    )

                    ######## SET SUBPLOT OPTIONS ######
                    # Set axis label and format for subplot
                    ax_to_plot.xaxis.set_major_locator(tick.MaxNLocator(integer=True))
                    ax_to_plot.tick_params(axis="both", which="major", labelsize=16)

                    # only put ylabel only on the first column subplots
                    if j == 0:
                        ax_to_plot.set_ylabel(
                            str(year) + " Dispatch (MWh)", fontsize=16
                        )

            ####### SET OVERALL FIGURE OPTIONS #######
            # Put legend on the far right of the plot
            ax_stor.legend(
                handles=legend_handles[::-1] + [ln_load] + [ln_loadst] + [ln_loadexp],
                fontsize=14,
                bbox_to_anchor=[leg_x, leg_y],
                loc="upper left",
                frameon=False,
                ncol=1,
            )

            # Set plot size (might have to play with the formula a bit)
            plot_length = len(self.days.value) * 5
            plot_height = len(self.years.value) * 3.5
            fig.set_size_inches(plot_length, plot_height)

            # for each representative day, add the title on top
            for i in range(len(self.days.value)):

                if len(self.days.value) > 1 and len(self.years.value) > 1:
                    ax[0, i].set_title("Day " + str(self.days.value[i]), fontsize=24)
                elif len(self.years.value) > 1:
                    ax[0].set_title("Day " + str(self.days.value[i]), fontsize=24)
                elif len(self.days.value) > 1:
                    ax[i].set_title("Day " + str(self.days.value[i]), fontsize=24)
                else:
                    ax.set_title("Day " + str(self.days.value[i]), fontsize=24)

            plt.show()

    def plot_heat_map(self, args=None):

        with self.output:
            self.output.clear_output()

            fig, ax = plt.subplots(figsize=(16, 10))

            # plotting the heatmap with custom format
            if self.hmap_col.value == "hourly_energy_cost_$_per_mwh":
                pivot = self.case.loads_and_power_balance.loc[
                    (self.case.loads_and_power_balance["period"] == self.year.value)
                    & self.case.loads_and_power_balance["zone"].isin(self.zones.value)
                ].pivot_table(
                    values=self.hmap_col.value,
                    index="day_id",
                    columns="hour_of_day",
                    aggfunc="mean",
                )

            elif self.hmap_col.value == "input_load_mw":
                pivot = self.case.loads_and_power_balance.loc[
                    (self.case.loads_and_power_balance["period"] == self.year.value)
                    & self.case.loads_and_power_balance["zone"].isin(self.zones.value)
                ].pivot_table(
                    values=self.hmap_col.value,
                    index="day_id",
                    columns="hour_of_day",
                    aggfunc="sum",
                )

            elif self.hmap_col.value in [
                "net_power_mw",
                "scheduled_curtailment_mw",
            ]:
                pivot = self.case.dispatch.loc[
                    (self.case.dispatch["period"] == self.year.value)
                    & self.case.dispatch["zone"].isin(self.zones.value)
                    & self.case.dispatch["resource_type"].isin(
                        self.resource_types.value
                    )
                ].pivot_table(
                    values=self.hmap_col.value,
                    index="day",
                    columns="hour_of_day",
                    aggfunc="sum",
                )

            else:
                pivot = self.case.transmit_power.loc[
                    (self.case.transmit_power["period"] == self.year.value)
                    & (
                        self.case.transmit_power["transmission_line"]
                        == self.tx_line.value
                    )
                ].pivot_table(
                    values=self.hmap_col.value,
                    index="day",
                    columns="hour_of_day",
                    aggfunc="mean",
                )

            if self.hmap_col.value == "hourly_energy_cost_$_per_mwh":
                heat_map = sns.heatmap(
                    pivot,
                    annot=True,
                    yticklabels=self.case.days,
                    cmap="coolwarm",
                    fmt=".0f",
                )

            else:
                heat_map = sns.heatmap(
                    pivot,
                    annot=False,
                    yticklabels=self.case.days,
                    cmap="coolwarm",
                    fmt=".0f",
                )

            plt.tick_params(axis="both", length=0, labelsize=10)
            plt.ylabel("RESOLVE Day", fontsize=14)
            plt.xlabel("Hour", fontsize=14)

            bottom, top = plt.ylim()

            plt.ylim(bottom + 0.5, top - 0.5)

            plt.show()
