import numpy as np
import matplotlib.pyplot as plt # for plotting

path = ['../results/Kuivajarvi2013_base/monthly/' , \
	'../results/Kuivajarvi2013_gravGill/monthly/', \
	'../results/Kuivajarvi2013_Goudsmit/monthly/', \
	'../results/Kuivajarvi2013_gravGillnoseiches/monthly/', \
	'../results/Kuivajarvi2013_Goudsmitnoseiches/monthly/', \
	'../results/Kuivajarvi2013_noseiches/monthly/'] # path to files
nheader = 12 #Number of lines in the header
year = '2013'
nameexp = ['base','GV+','IS+','GV+SS-','IS+SS-','SS-'] #,'nobotdrag','noseiches']

nm = input("Enter month \n")
#nd = input("Enter day \n")
#nhstep = input("Enter time step, hours \n")
#nh = 0

#varname = 'TKE'
varname = 'TKE_balance_terms'
units = '$m^2/s^3$'
xlabel = varname + ', ' + units
ylabel = 'Depth, m'

#ncol = [7, 8, 9, 10,11] #The column number to display, 2 for temperature
ncol = [8, 9, 10] #The column number to display, 2 for temperature
#ncol = [6]
#ncol = [11]
ncols = 12 #Total number of columns
nvars = len(ncol)#Number of variables to plot

nfiles = len(path) #number of files
time = [year + str(nm).zfill(2)]

#legitems = ['TKE dissipation', 'TKE buoyancy production', 'TKE shear production', \
#	    'TKE production by seiches', 'TKE turbulent transport']
legitems = ['$B$', '$S$', '$S_{seiche}$']
#legitems = ['TKE']
#legend = ['Heat conductance, $m^2/s$']

#Set of types and colors of lines
colorline = ['b','g','r','c','m','y','k']
typeline = ['-','--','-.','v','o','^','.',':',',','1']
ncolors = len(colorline)

#Files names suffixes (timestamps)
#while nh + nhstep <= 23:
#	 nfiles += 1
#	 nh += nhstep # Incrementing by nhstep
#	 time.append(year + str(nm).zfill(2) + str(nd).zfill(2) + str(nh).zfill(2))
#	 legend.append(str(nh).zfill(2) + 'h')
legend = []
for i in range(nfiles):
	for j in range(nvars):
		legend.append(legitems[j] + '_' + nameexp[i])
	fid = open(path[i]+'Profiles  1  1'+str(time[0])+'.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])#depths are given by negative values in the input file
		for k in range(nvars):
			vals[j,k] = float(line[ncol[k]-1])
	for j in range(nvars):
		jcolor = j-(j//ncolors)*ncolors # periodic use of colors in one experiment
		#plt.plot(vals[:,j],depths,colorline[jcolor]+typeline[i],linewidth=2)
		plt.semilogx(vals[:,j],depths,colorline[jcolor]+typeline[i],linewidth=2)
#print(depths)
#print(vals)

plt.gca().invert_yaxis()
plt.xlabel(xlabel,fontsize=16)
plt.ylabel(ylabel,fontsize=16)
#plt.legend(legend,loc='best',fontsize=16)
plt.legend(legend,bbox_to_anchor=(0.95, 0.95), loc=2, borderaxespad=0.,fontsize=15)
plt.tick_params(labelsize=16)

#fig = plt.gcf()
#fig.set_size_inches(18.5, 10.5)

# Filename postfix
postfix = nameexp[0]
for i in range(nfiles-1):
	postfix = postfix + '&' + nameexp[i+1]

#Saving file
for i in range(nfiles):
	plt.savefig(path[i] + varname + '_' + time[0] + '_' + postfix +'_monthly.png')

plt.show()
