import os
import pickle
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

matplotlib.rc('text', usetex=True)
pd.options.display.notebook_repr_html = False


def plot_optimization_result(datasource, directory):  # data source is dataframe
    sns.set_theme(style='whitegrid')  # "white", "dark", "whitegrid", "darkgrid", "ticks"
    plt.rcParams["figure.figsize"] = (16, 9)
    fig, axs = plt.subplots(2, 2)
    plt.subplots_adjust(wspace=0.7, hspace=0.3)
    plt.autoscale(tight=True)
    T = np.array([i for i in range(24)])

    # plot step cost in ax[0]
    axs[0, 0].cla()
    axs[0, 0].set_ylabel('Costs')
    axs[0, 0].set_xlabel('Time(h)')
    axs[0, 0].bar(T, datasource['step_cost'])
    # axs[0,0].set_xticks([i for i in range(24)],[i for i in range(1,25)])

    # plot soc and price in ax[1]
    axs[0, 1].cla()
    axs[0, 1].set_ylabel('Price')
    axs[0, 1].set_xlabel('Time(h)')

    axs[0, 1].plot(T, datasource['price'], drawstyle='steps-mid', label='Price', color='pink')
    axs[0, 1] = axs[0, 1].twinx()

    axs[0, 1].set_ylabel('SOC')
    axs[0, 1].plot(T, datasource['soc'], drawstyle='steps-mid', label='SOC', color='grey')
    # axs[0,1].set_xticks([i for i in range(24)],[i for i in range(1,25)])
    axs[0, 1].legend(loc='upper right', fontsize=12, frameon=False, labelspacing=0.3)

    # plot accumulated generation and consumption in ax[2]
    axs[1, 0].cla()
    axs[1, 0].set_ylabel('Outputs of DGs and Battery')
    axs[1, 0].set_xlabel('Time(h)')
    battery_positive = np.array(datasource['battery_energy_change'])
    battery_negative = np.array(datasource['battery_energy_change'])
    battery_negative = np.minimum(battery_negative, 0)  # discharge
    battery_positive = np.maximum(battery_positive, 0)  # charge
    # deal with power exchange within the figure
    imported_from_grid = np.array(datasource['grid_import'])
    exported_2_grid = np.array(datasource['grid_export'])
    axs[1, 0].bar(T, datasource['gen1'], label='Gen1')
    axs[1, 0].bar(T, datasource['gen2'], label='Gen2', bottom=datasource['gen1'])
    axs[1, 0].bar(T, datasource['gen3'], label='Gen3', bottom=datasource['gen2'] + datasource['gen1'])
    axs[1, 0].bar(T, -battery_positive, color='blue', hatch='/', label='ESS charge')
    axs[1, 0].bar(T, -battery_negative, hatch='/', label='ESS discharge',
                  bottom=datasource['gen3'] + datasource['gen2'] + datasource['gen1'])
    # import as generate
    axs[1, 0].bar(T, imported_from_grid, label='Grid import',
                  bottom=-battery_negative + datasource['gen3'] + datasource['gen2'] + datasource['gen1'])
    # export as load
    axs[1, 0].bar(T, -exported_2_grid, label='Grid export', bottom=-battery_positive)
    axs[1, 0].plot(T, datasource['netload'], label='Netload', drawstyle='steps-mid', alpha=0.7)
    axs[1, 0].legend(loc='upper right', fontsize=12, frameon=False, labelspacing=0.3)
    # axs[1,0].set_xticks([i for i in range(24)],[i for i in range(1,25)])

    fig.savefig(f"{directory}/optimization_information.svg", format='svg', dpi=600, bbox_inches='tight')
    print('optimization results have been ploted and saved')


def plot_evaluation_information(datasource, directory):
    sns.set_theme(style='whitegrid')
    with open(datasource, 'rb') as tf:
        test_data = pickle.load(tf)
    # plot unbalance, and reward of each step by bar figures
    plt.rcParams["figure.figsize"] = (16, 9)
    fig, axs = plt.subplots(2, 2)
    plt.subplots_adjust(wspace=0.7, hspace=0.3)
    plt.autoscale(tight=True)

    # prepare data for evaluation the environment here
    eval_data = pd.DataFrame(test_data['system_info'])
    eval_data.columns = ['time_step', 'price', 'netload', 'action', 'real_action', 'soc', 'battery', 'gen1', 'gen2',
                         'gen3', 'temperature', 'irradiance', 'unbalance', 'operation_cost']

    # plot unbalance in axs[0]
    axs[0, 0].cla()
    axs[0, 0].set_ylabel('Unbalance of Generation and Load')
    axs[0, 0].bar(eval_data['time_step'], eval_data['unbalance'], label='Exchange with Grid', width=0.4)
    axs[0, 0].bar(eval_data['time_step'] + 0.4, eval_data['netload'], label='Netload', width=0.4)
    axs[0, 0].legend(loc='upper right', fontsize=12, frameon=False, labelspacing=0.5)
    # axs[0,0].set_xticks([i for i in range(24)],[i for i in range(1,25)])

    # plot energy charge/discharge with price in ax[1]
    axs[0, 1].cla()
    axs[0, 1].set_ylabel('Price')
    axs[0, 1].set_xlabel('Time Steps')

    axs[0, 1].plot(eval_data['time_step'], eval_data['price'], drawstyle='steps-mid', label='Price', color='pink')
    axs[0, 1] = axs[0, 1].twinx()
    axs[0, 1].set_ylabel('SOC')
    # axs[0,1].set_xticks([i for i in range(24)], [i for i in range(1, 25)])
    axs[0, 1].plot(eval_data['time_step'], eval_data['soc'], drawstyle='steps-mid', label='SOC', color='grey')
    axs[0, 1].legend(loc='upper right', fontsize=12, frameon=False, labelspacing=0.3)

    # plot generation and netload in ax[2]
    axs[1, 0].cla()
    axs[1, 0].set_ylabel('Outputs of Units and Netload (kWh)')
    # axs[1,0].set_xticks([i for i in range(24)], [i for i in range(1, 25)])
    battery_positive = np.array(eval_data['battery'])
    battery_negative = np.array(eval_data['battery'])
    battery_positive = np.maximum(battery_positive, 0)  # charge
    battery_negative = np.minimum(battery_negative, 0)  # discharge

    # deal with power exchange within the figure 
    imported_from_grid = np.minimum(np.array(eval_data['unbalance']), 0)
    exported_2_grid = np.maximum(np.array(eval_data['unbalance']), 0)
    x = eval_data['time_step']
    axs[1, 0].bar(x, eval_data['gen1'], label='Gen1')
    axs[1, 0].bar(x, eval_data['gen2'], label='Gen2', bottom=eval_data['gen1'])
    axs[1, 0].bar(x, eval_data['gen3'], label='Gen3', bottom=eval_data['gen1'] + eval_data['gen2'])
    axs[1, 0].bar(x, -battery_positive, color='blue', hatch='/', label='ESS charge')
    axs[1, 0].bar(x, -battery_negative, label='ESS discharge', hatch='/',
                  bottom=eval_data['gen1'] + eval_data['gen2'] + eval_data['gen3'])
    axs[1, 0].bar(x, -imported_from_grid, label='Grid import',
                  bottom=eval_data['gen1'] + eval_data['gen2'] + eval_data['gen3'] - battery_negative)
    axs[1, 0].bar(x, -exported_2_grid, label='Grid export', bottom=-battery_positive)

    axs[1, 0].plot(x, eval_data['netload'], drawstyle='steps-mid', label='Netload')
    axs[1, 0].legend(loc='upper right', fontsize=12, frameon=False, labelspacing=0.3)

    # plot reward in axs[3]
    axs[1, 1].cla()
    axs[1, 1].set_ylabel('Costs')
    axs[1, 1].bar(eval_data['time_step'], eval_data['operation_cost'])
    fig.savefig(f"{directory}/evaluation_information.svg", format='svg', dpi=600, bbox_inches='tight')
    print('evaluation figure have been plot and saved')


def make_dir(directory, feature_change):
    cwd = f'{directory}/DRL_{feature_change}_plots'
    os.makedirs(cwd, exist_ok=True)
    return cwd


class PlotArgs:
    def __init__(self) -> None:
        self.cwd = None
        self.feature_change = None
        self.plot_on = None