# -*- coding: utf-8 -*-
"""
"""
from config import clustered_path, original_path, plot_path

from os import path, listdir

import matplotlib
from matplotlib import pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np
import pandas as pd

def plot_heatmap(buses, difference):
    """
    """
    x = np.array(buses['x'])
    y = np.array(buses['y'])

    alpha = np.array([i for i in difference.values()])

    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(12, 8)
    cmap = plt.cm.jet

    plt.hexbin(x, y, C=alpha, cmap=cmap, gridsize=20)
    cb = plt.colorbar()

    cb.set_label('Absolute difference of Installed Storage Capacities in MW')
    return(fig)


# TODO: check why are there multiple storages for e.g. bus with name 2 and 9

original = pd.read_csv(path.join(original_path, 'storage_units.csv'))
buses = pd.read_csv(path.join(original_path, 'buses.csv'))

# store all storage results
clustered = {}
# nested dict with absolute difference per cluster and bus
abs_diff = {}

c_df = pd.DataFrame()
with PdfPages(path.join(plot_path, 'storage-capacities.pdf')) as pdf:
    for d in listdir(clustered_path):
        # TODO: excluse distance matrix, move this file to the root directory?
        if d != 'Z.csv':
            abs_diff[d] = {}
            df = pd.read_csv(path.join(clustered_path, d, 'storage_units.csv'))
            df['cdays'] = int(d)
            df['abs_diff'] = df.p_nom_opt - original.p_nom_opt
            abs_diff[d] = {
                b: sum(df.loc[df['bus'] == b]['abs_diff'])
                for b in buses['name']}
            c_df = pd.concat([c_df, df])
            clustered[d] = df
            fig = plot_heatmap(buses, abs_diff[d])
            pdf.savefig(fig)


# experimental stuff
c_df.set_index(['cdays', 'name'], inplace=True)
c_df.sort_index()

@@ -298,16 +298,36 @@
 fference
+, minmax=(None,None)
 ):%0A    %22
@@ -595,34 +595,126 @@
 e=20
-)%0A    cb = plt.colorbar()%0A
+, vmin=minmax%5B0%5D,%0A               vmax=minmax%5B1%5D)%0A%0A    cb = plt.colorbar()%0A%0A    # plt.clim(-100, 100) normalized colors
 %0A   
@@ -1157,82 +1157,8 @@
 e()%0A
-with PdfPages(path.join(plot_path, 'storage-capacities.pdf')) as pdf:%0A    
 for 
@@ -1187,20 +1187,16 @@
 _path):%0A
-    
     # TO
@@ -1266,20 +1266,16 @@
 ry?%0A    
-  
-  
 if d != 
@@ -1283,36 +1283,32 @@
 Z.csv':%0A        
-   
- 
 abs_diff%5Bd%5D = %7B%7D
@@ -1308,20 +1308,16 @@
 d%5D = %7B%7D%0A
-    
         
@@ -1388,28 +1388,24 @@
 '))%0A        
-   
- 
 df%5B'cdays'%5D 
@@ -1413,20 +1413,16 @@
  int(d)%0A
-    
         
@@ -1476,28 +1476,24 @@
 opt%0A        
-    
 abs_diff%5Bd%5D 
@@ -1496,20 +1496,16 @@
 %5Bd%5D = %7B%0A
-    
         
@@ -1563,20 +1563,16 @@
         
-    
 for b in
@@ -1587,20 +1587,16 @@
 name'%5D%7D%0A
-    
         
@@ -1632,20 +1632,16 @@
         
-    
 clustere
@@ -1649,16 +1649,306 @@
 %5Bd%5D = df
+%0A%0A# plot%0Aminmax = (min(%5Bmin(%5Bj for j in i.values()%5D) for i in abs_diff.values()%5D),%0A          max(%5Bmax(%5Bj for j in i.values()%5D) for i in abs_diff.values()%5D))%0A%0Awith PdfPages(path.join(plot_path, 'storage-capacities.pdf')) as pdf:%0A    for d in listdir(clustered_path):%0A        if d != 'Z.csv':
 %0A       
@@ -1989,16 +1989,31 @@
 _diff%5Bd%5D
+, minmax=minmax
 )%0A      
@@ -2036,17 +2036,16 @@
 g(fig)%0A%0A
-%0A
 # experi
@@ -2057,16 +2057,18 @@
 l stuff%0A
+# 
 c_df.set
@@ -2107,16 +2107,18 @@
 e=True)%0A
+# 
 c_df.sor

Make colors consitent for one result set