import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.ticker import FuncFormatter
import numpy as np
import pandas as pd
import os
import seaborn as sns
import json
from matplotlib.colors import LinearSegmentedColormap

def read_data(file_name: str):
    with open(file_name, 'r') as f:
        data = json.load(f)
    for key, value in data.items():
        for subkey, subvalue in value.items():
            data[key][subkey] = float(subvalue)
    df = pd.DataFrame.from_dict(data, orient='index')
    df = df.T
    df.index = pd.to_numeric(df.index)
    df.columns = pd.to_numeric(df.columns)
    return df

def draw_results(results, filename, title_benefit, annot_benefit=False, figure_size=(10, 10)):
    df=results
    df = df.astype(float)
    df.index = df.index / 1000
    df.index = df.index.map(int)
    df.columns = df.columns / 1000
    df.columns = df.columns.map(int)
    min_value = df.min().min()
    max_value = df.max().max()
    max_scale = max(abs(min_value/1000), abs(max_value/1000))

    df[df.columns[-1] + 1] = df.iloc[:, -1] 
    new_Data = pd.DataFrame(index=[df.index[-1] + 1], columns=df.columns)
    for i in df.columns:
        new_Data[i] = df[i].iloc[-1]
    # print(new_Data)
    df = pd.concat([df, new_Data])

    X, Y = np.meshgrid(np.arange(df.shape[1]), np.arange(df.shape[0]))

    def fmt(x,pos):
        return '{:.0f}'.format(x/1000)

    cmap = sns.color_palette("coolwarm", as_cmap=True)
    plt.figure(figsize=figure_size)
    ax = sns.heatmap(df/1000, fmt=".1f", cmap=cmap, vmin=-max_scale, vmax=max_scale, annot=annot_benefit)
    CS = ax.contour(X, Y, df,  colors='black', alpha=0.5)
    ax.clabel(CS, inline=True, fontsize=10, fmt=FuncFormatter(fmt))
    plt.title(title_benefit)
    plt.gca().invert_yaxis()
    plt.xlim(0, df.shape[1] - 1)
    plt.ylim(0, df.shape[0] - 1)
    plt.xlabel('ESS Capacity (MWh)')
    plt.ylabel('PV Capacity (MW)')
    plt.savefig(filename)

def draw_cost(costs, filename, title_cost, annot_cost=False, figure_size=(10, 10)):
    df = costs
    df = df.astype(int)
    df.index = df.index / 1000
    df.index = df.index.map(int)
    df.columns = df.columns / 1000
    df.columns = df.columns.map(int)

    df[df.columns[-1] + 1] = df.iloc[:, -1] 
    new_Data = pd.DataFrame(index=[df.index[-1] + 1], columns=df.columns)
    for i in df.columns:
        new_Data[i] = df[i].iloc[-1]
    # print(new_Data)
    df = pd.concat([df, new_Data])
    X, Y = np.meshgrid(np.arange(df.shape[1]), np.arange(df.shape[0]))

    def fmt(x, pos):
        return '{:.0f}'.format(x / 1000000)

    plt.figure(figsize=figure_size)
    ax = sns.heatmap(df/1000000,  fmt=".1f", cmap='viridis', annot=annot_cost)
    CS = ax.contour(X, Y, df,  colors='black', alpha=0.5)
    ax.clabel(CS, inline=True, fontsize=10, fmt=FuncFormatter(fmt))
    plt.title(title_cost)
    plt.gca().invert_yaxis()
    plt.xlim(0, df.shape[1] - 1)
    plt.ylim(0, df.shape[0] - 1)
    plt.xlabel('ESS Capacity (MWh)')
    plt.ylabel('PV Capacity (MW)')
    plt.savefig(filename)


def draw_overload(overload_cnt, filename, title_unmet, annot_unmet=False, figure_size=(10, 10)):
    df = overload_cnt
    df = (4 * 24 * 365 - df) / (4 * 24 * 365)
    df = df.astype(float)
    df.index = df.index / 1000
    df.index = df.index.map(int)
    df.columns = df.columns / 1000
    df.columns = df.columns.map(int)
    print('unmet index')
    print(df.index)
    print('unmet columns')
    print(df.columns)
    min_value = df.min().min()
    max_value = df.max().max()
    # max_scale = max(abs(min_value/1000), abs(max_value/1000))


    df[df.columns[-1] + 1] = df.iloc[:, -1] 
    new_Data = pd.DataFrame(index=[df.index[-1] + 1], columns=df.columns)
    for i in df.columns:
        new_Data[i] = df[i].iloc[-1]
    # print(new_Data)
    df = pd.concat([df, new_Data])


    plt.figure(figsize=figure_size)
    cmap = LinearSegmentedColormap.from_list("", ["white", "blue"])
    ax = sns.heatmap(df, fmt=".00%", cmap=cmap, vmin=0, vmax=1, annot=annot_unmet)

    cbar = ax.collections[0].colorbar
    cbar.set_ticks([0, 0.25, 0.5, 0.75, 1])
    cbar.set_ticklabels(['0%', '25%', '50%', '75%', '100%'])
    cbar.ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: f'{x:.0%}'))
    X, Y = np.meshgrid(np.arange(df.shape[1]), np.arange(df.shape[0]))

    def fmt(x, pos):
        return '{:.0f}%'.format(x * 100)
    CS = ax.contour(X, Y, df,  colors='black', alpha=0.5)

    ax.clabel(CS, inline=True, fontsize=10, fmt=FuncFormatter(fmt))

    plt.xlim(0, df.shape[1] - 1)
    plt.ylim(0, df.shape[0] - 1)
    plt.title(title_unmet)
    plt.xlabel('ESS Capacity (MWh)')
    plt.ylabel('PV Capacity (MW)')
    plt.savefig(filename)


directory = 'data/'

file_list = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]


split_files = [f.split('-') for f in file_list]

for f in split_files:
    print(f[-1])
costs_files = [f for f in split_files if f[-1].endswith('costs.json')]
print(f'find costs files: {costs_files}')
overload_files = [f for f in split_files if f[-1].endswith('overload_cnt.json')]
print(f'find overload files: {overload_files}')
results_files = [f for f in split_files if f[-1].endswith('results.json')]
print(f'find results files: {results_files}')

costs_dfs = [read_data(directory + '-'.join(f)) for f in costs_files]
overload_dfs = [read_data(directory + '-'.join(f)) for f in overload_files]
results_dfs = [read_data(directory + '-'.join(f)) for f in results_files]

for costs_df, overload_df, results_df in zip(costs_dfs, overload_dfs, results_dfs):

    draw_cost(costs_df, 
              f'plots/costs-{int(costs_df.columns[-1])}-pv-{int(costs_df.index[-1])}.png', 
              f'Costs for ESS-{int(costs_df.columns[-1])}-pv-{int(costs_df.index[-1])}MWh', annot_cost=False)

    draw_overload(overload_df, 
                  f'plots/overload-ess-{overload_df.columns[-1]}-pv-{overload_df.index[-1]}.png', 
                  f'Overload for ess-{overload_df.columns[-1]}MW pv-{overload_df.index[-1]}MWh', 
                  annot_unmet=False)

    draw_results(results_df, 
                 f'plots/results-{results_df.columns[-1]}-pv-{results_df.index[-1]}.png', 
                 f'Results for ess-{results_df.columns[-1]}MW pv-{results_df.index[-1]}MWh', annot_benefit=False)