import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt # for plotting
from mpl_toolkits.axes_grid1 import host_subplot
import mpl_toolkits.axisartist as AA

mpl.rc('mathtext', default='regular')

path = '../results/BolshoiVilui20153_ext05/daily/' # path to files
nheader = 7 #Number of lines in the header
year = '2015'

nm = input("Enter month \n")
nd = input("Enter first day \n")
ndl = input("Enter last day \n")
nhstep = input("Enter time step, days \n")
nh = nd - nhstep
#ndays = 30

#varname = ['TKE']
#units = ['$m^2/s^2$']
#ncol = [4]
varname = ['Temperature','Salinity','Eddy heat conductivity']
units = ['$C^{\circ}$','kg/kg','$m^2/s$']
ncol = [2,3,6]
nvars = len(varname) #Number of variables to plot
legend = [[] for _ in range(nvars)] # Creating list of different lists
plotid = [[] for _ in range(nvars)] # Creating list of different lists
ax = [None]*nvars
#varname = 'Dissipation rate'
#varname = 'Heat conductance'

#Tick label size on plots
mpl.rc('xtick', labelsize=20)
mpl.rc('ytick', labelsize=20)

mpl.rc('font', size=20)
#mpl.rc('ylabel', labelsize=20)

xlabel = []
for i in range(nvars):
        xlabel.append(varname[i] + ', ' + units[i])
ylabel = 'Depth, m'

ncols = 6 #Total number of columns

nfiles = 0 #number of files
time = []
#legend = ['TKE']
#legend = []
#Files names suffixes (timestamps)
print(legend)
while nh + nhstep <= ndl: 
        nfiles += 1
        nh += nhstep # Incrementing by nhstep
        time.append(year + str(nm).zfill(2) + str(nh).zfill(2))
        #Creating a legend for each variable
        for i in range(nvars):
                #legend[i].append(varname[i] + ' ' + str(nh).zfill(2) + ' day')
                legend[i].append(varname[i])

for i in range(nfiles):
        fid = open(path+'Profiles  1  1'+str(time[i])+'.dat','r')
        #Reading the header
        for j in range(nheader):
	        line = fid.readline()
        #if i == 0 :
        #Counting the number of raws in a file (assuming the same in all files)
        array = fid.read()
        nraws = array.count('\n')
        print('The number of levels is '+str(nraws))
        depths = [0]*nraws
        vals = np.zeros((nraws,nvars))
        fid.seek(0) #rewinding file
        for j in range(nheader):
        	line = fid.readline()
        # Reading array of depths and values to be plotted
        for j in range(nraws) :
        	line = fid.readline()
        	line = line.split()
        	depths[j] = - float(line[0])
                for k in range(nvars):
	                vals[j,k] = float(line[ncol[k]-1])
        for j in range(nvars):
                if i == 0 and j == 0:
                        ax[j] = host_subplot(111,axes_class=AA.Axes) #plt.subplots()
                        ax[j].plot(vals[:,j],depths,linewidth=2)
                        #Moving the top boundary to get space for second axis
                        pos0 = ax[j].get_position()
                        pos1 = [pos0.x0, pos0.y0,  pos0.width, pos0.height - 0.1]
                        ax[j].set_position(pos1)
                else:
                        if j > 0 and ax[j] == None:
                                ax[j] = ax[0].twiny()
                        ax[j].plot(vals[:,j],depths,linewidth=2)
#print(vals)

plt.gca().invert_yaxis()
plt.ylabel(ylabel,fontsize=30)
for i in range(nvars):
        #Moving axis in vertical
        if i > 1:
                offset = 50
                new_fixed_axis = ax[i].get_grid_helper().new_fixed_axis
                ax[i].axis["top"] = new_fixed_axis(loc="top",axes=ax[i],offset=(0, offset))
                ax[i].axis["top"].toggle(all=True)
        ax[i].legend(legend[i],loc=i+1,fontsize=20)
        ax[i].set_xlabel(xlabel[i],fontsize=30)
        #plt.ticklabel_format(style='sci', axis='x')
	#params = {'axes.labelsize': 18,'axes.titlesize':20,'xtick.labelsize': 20, 'ytick.labelsize': 20}
	#mpl.rcParams.update(params)
	#ax[i].xaxis.label.set_size(20)

plt.savefig(path + varname[0] + '_' + year + str(nm).zfill(2) + '_daily.png')
plt.show()
