'''
Created on 20170511
@author: Cherry
'''
import csv
from matplotlib import cm
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from matplotlib.ticker import FormatStrFormatter
from matplotlib.pyplot import savefig



degrees = []
kshells = []
upper_values = []
eu_values = []
el_values = []
lower_values = []

root = 'D:/Cherry/data/netsci'
dataPath = root + '/result_f_distribution.csv'

def readData():
    with open(dataPath, 'r') as f:
        i = 0
        for row in csv.reader(f.read().splitlines()):
            if i == 0:
                i += 1
            else:
                id, kshell, degree, upper, eu, el, lower = [i for i in row]
                degrees.append(degree)
                kshells.append(kshell)
                upper_values.append(upper)
                eu_values.append(eu)
                el_values.append(el)
                lower_values.append(lower)
            
def getSize():
    max_degree = 0
    max_kshell = 0
    for degree in degrees:
        if max_degree < int(degree):
            max_degree = int(degree)
            
    for kshell in kshells:
        if max_kshell < kshell:
            max_kshell = kshell
    return max_degree, max_kshell

def generateZ(max_degree, max_kshell):
    max_degree = int(max_degree)
    max_kshell = int(max_kshell)
    z_upper_values = [None] * max_kshell
    z_eu_values = [None] * max_kshell
    z_el_values = [None] * max_kshell
    z_lower_values = [None] * max_kshell
    for i in range(len(z_upper_values)):
        z_upper_values[i] = [0] * max_degree
        z_eu_values[i] = [0] * max_degree
        z_el_values[i] = [0] * max_degree
        z_lower_values[i] = [0] * max_degree
        
        
    for j in range(len(degrees)):
        index_x = degrees[j]  # get the degree named index_x, then regard the index_x-1 as the zvalue[][]'s second dimention index.   zvalue[][index-x-1]
        index_y = kshells[j]  # get the kshell named index_y, then regard the max-kshell-index_y as the zvalue[][]'s first dimention index.  zvalue[max_kshell][]
        upperValue = upper_values[j]
        z_upper_values[int(max_kshell) - int(index_y)][int(index_x) - 1] = float(upperValue)
        euValue = eu_values[j]
        z_eu_values[int(max_kshell) - int(index_y)][int(index_x) - 1] = float(euValue)
        elValue = el_values[j]
        z_el_values[int(max_kshell) - int(index_y)][int(index_x) - 1] = float(elValue)
        lowerValue = lower_values[j]
        z_lower_values[int(max_kshell) - int(index_y)][int(index_x) - 1] = float(lowerValue)
        
    return z_upper_values, z_eu_values, z_el_values, z_lower_values
        
def draw_heatmap(data0, xlabels, ylabels, savepath, titleStr):
    cmap = cm.Blues
    figure = plt.figure(facecolor='w')
    
    # set the font
    font = {'family' : 'serif', 'color'  : 'darkred', 'weight' : 'normal', 'size'   : 16, }
    titlefont = {'family' : 'serif', 'color'  : 'black', 'weight' : 'normal', 'size'   : 20, }
    
    ax = figure.add_subplot(1, 1, 1, position=[0.8, 0.15, 1, 1])
    ax.set_yticks(range(len(ylabels)))
    ax.set_yticklabels(ylabels)
    ax.set_xticks(range(len(xlabels)))
    ax.set_xticklabels(xlabels)
    ax.set_ylabel('kshell', fontdict=font)  # set labelY to show the meaning of axiX
    ax.set_xlabel('degree', fontdict=font)  # set labelX to show the meaning of axiY
    plt.title(titleStr, fontdict=titlefont)  # set title to show the meaning of the distribution graph
    
    xmajorLocator = MultipleLocator(5)  # set the main scale of axiX to the multiple of 5
    xmajorFormatter = FormatStrFormatter('%u')  # set the famat of the labelX
#     xminorLocator   = MultipleLocator(5) 
    ax.xaxis.set_major_locator(xmajorLocator)
    ax.xaxis.set_major_formatter(xmajorFormatter)
    vmax = data0[0][0]
    vmin = data0[0][0]
    for i in data0:
        for j in i:
            if j > vmax:
                vmax = j
            if j < vmin:
                vmin = j
    map0 = ax.imshow(data0, interpolation='nearest', cmap=cmap, aspect='auto', vmin=vmin, vmax=vmax)
    plt.colorbar(mappable=map0, cax=None, ax=None, shrink=0.5)
    # plt.show()
    savefig(savepath)
    plt.clf()
    
      
readData()
max_degree, max_kshell = getSize()
z_upper_values, z_eu_values, z_el_values, z_lower_values = generateZ(max_degree, max_kshell)
x_degrees = []
y_kshells = []
i = 1
while(i <= int(max_degree)):  # the value at index0 is the min value
    x_degrees.append(int(i))
    i += 1
    
j = int(max_kshell)
while(j > 0):  # the value at index0 is the max value
    y_kshells.append(int(j))
    j -= 1

draw_heatmap(z_upper_values, x_degrees, y_kshells, root + '/distri_upper.png', 'upper')
draw_heatmap(z_eu_values, x_degrees, y_kshells, root + '/distri_eu.png', 'eu')
draw_heatmap(z_el_values, x_degrees, y_kshells, root + '/distri_el.png', 'el')
draw_heatmap(z_lower_values, x_degrees, y_kshells, root + '/distri_lower.png', 'lower')
print('finished')