import seaborn as sns
import pandas as pd
import matplotlib.colors as mc

# For defining color map
def NonLinCdict(steps, hexcol_array):
    cdict = {'red': (), 'green': (), 'blue': ()}
    for s, hexcol in zip(steps, hexcol_array):
        rgb = mc.hex2color(hexcol)
        cdict['red'] = cdict['red'] + ((s, rgb[0], rgb[0]),)
        cdict['green'] = cdict['green'] + ((s, rgb[1], rgb[1]),)
        cdict['blue'] = cdict['blue'] + ((s, rgb[2], rgb[2]),)
    return cdict

# Load data
thedata = pd.read_csv("data/month_hour_pct_matrix.csv")
mns = thedata.iloc[0:12, 0]
pmat = thedata.iloc[0:12, 1:25]

# Set figure size
sns.set(rc={'figure.figsize':(11,4.5), 'pdf.fonttype': 42})

# Make color map
# hc = ['#dff5ee', '#92dee9', '#3fc1f9', '#5598fd', '#a75bd8', '#d4007f']
hc = ['#f0f0f0', '#c3dde3', '#acc6cc', '#aeaaaa', '#c0877c', '#da5241']
th = [0, 0.1, 0.2, 0.3, .4, 1]
cdict = NonLinCdict(th, hc)
cm = mc.LinearSegmentedColormap('custom', cdict)

# Heatmap
p1 = sns.heatmap(
    pmat,
    square=True,
    linewidths=0.1,
    yticklabels=mns,
    cmap=cm)
p1.set_xticks([0.5,6.5,12.5,18.5,23.5])
p1.set_xticklabels(["12am", "6am", "noon", "6pm", "11pm"])
fig = p1.get_figure()

# Save file
fig.savefig("pmat-heatmap.pdf")