##############################
# Heatmap from the VAR files #
##############################


# Import modules
import pencil as pc
import seaborn as sns 
import matplotlib.pyplot as plt 
import pandas as pd 
import numpy as np 
  
# create figure environment:
f, ax = plt.subplots(figsize=(10,5))

# maximum number of VAR files:
numberofvarfiles = len(open('../data/proc0/varN.list').readlines(  ))-1

# collect all density snapshots:
lnrho = []
for i in range(0, numberofvarfiles, 1):
    var = pc.read.var(ivar=i,trimall=True, datadir='../data')
    if i==0:
       lnrho = var.lnrho[0,0,:]
    else:
       lnrho = np.vstack((lnrho, var.lnrho[0,0,:]))
  
# create heat map:
sns.heatmap(np.transpose(lnrho), ax=ax)


# add labels
ax.collections[0].colorbar.set_label("$\\mathrm{log}(\\rho)$")
plt.ylabel("$x$ axis")
plt.xlabel("number of VAR file")

f.savefig("snapshots_heatmap.pdf" , bbox_inches='tight')