import numpy as nm
import matplotlib.pyplot as plt # for plotting
#import os
import time as tm
import os
import math as mt
from plot_timser import *
import matplotlib.dates as dts
import datetime as dtm
from matplotlib import ticker

nplot = input("Choose a variable to plot \n \
 1 - temperature \n \
 2 - salinity \n \
 3 - density (piece-wise constant) \n \
 4 - CH_4 \n \
 5 - O_2 \n \
 6 - CO_2 \n \
 7 - CH_4 bubble flux \n \
 8 - Richardson number \n \
 9 - Rp number \n \
 10 - Rpdens number \n \
 11 - x-speed \n \
 12 - y-speed \n \
 13 - Heat conductance \n \
 14 - TKE dissipation rate \n \
 15 - temperature measured \n \
 16 - oxygen measured \n \
 17 - methane measured \n \
 18 - carbon dioxide measured \n ")

nmod_max = 12 #The last item in the list above, from the modeled data

path = '../results/Kuiva2014septnew/time_series/' # path to files
#pathobs = os.path.expanduser('/media/victor/main/Files/main/observdata/Iseo/')
pathobs = os.path.expanduser('/media/victor/main/Files/main/observdata/Kuivajarvi/Kuiva2014sept_campaign/')
exp = 'base'

#nlevs_water = 16
nlevplot = 20 # Number of countour levels

#Year, month, ... for the 1-st entry of the time series
ny = 2016
mdays = 365/12
nmonth = 9
nday = 2
nhour = 0
nmin = 0
timezone = None

logscalebar = False

# Choose for date-time format of labels on X axis
time_label_date = True#False

if nplot >= 1: #and nplot <= nmod_max:
        time0_date = dtm.datetime(ny,nmonth,nday,nhour,nmin,0,0,timezone)
        time0_ = dts.date2num(time0_date)
        years  = dts.YearLocator()   # every year
        months = dts.MonthLocator()  # every month
        days   = dts.DayLocator()  # every day
        dateFmt = dts.DateFormatter('%y-%m-%d')

# User controls
if int(nplot) == 1: # Temperature
        basename = 'water_temp  1  1'
        fileout = 'temp_wat_cont'
        nheader = 6
        labelx = 'Time' #, months
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time of data
        time0 = 0. #month_doy(7,ny)/mdays + 17./31.
        #t0disp = 7.50 # Initial time for display
        #t1disp = 7.95 # Final time for display
        #nvardisp = nlevs_water
        title = "Temperature, $^\circ$C "
        levmin = 8.
        levmax = 20.
        nlevplot = 20 # Number of countour levels
        pathsave = path
if int(nplot) == 2: # Salinity
        basename = 'sal_water  1  1'
        fileout = 'sal_wat_cont'
        nheader = 6
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 7.48 #Initial time of data
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 7.50 # Initial time for display
        #t1disp = 7.95 # Final time for display
        #nvardisp = nlevs_water
        title = "Salinity, $kg/kg$ "
        levmin = 0.E-3
        levmax = 32.E-3
        nlevplot = 16 # Number of countour levels
        pathsave = path
	stepticks = 4.E-3
if int(nplot) == 3: # Density piece-wise constant
        basename = 'water_dens_layers  1  1'
        fileout = 'dens_wat_cont'
        nheader = 6
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 7.48 #Initial time of data
        time0 = 0. #month_doy(7,ny)/mdays + 17./31.
        #t0disp = 7.50 # Initial time for display
        #t1disp = 7.95 # Final time for display
        #nvardisp = nlevs_water
        title = "Density, kg/m$^3$ "
        #levmin = 0.E-3
        #levmax = 32.E-3
        nlevplot = 40 # Number of countour levels
        pathsave = path
elif int(nplot) == 4: # Methane
        basename = 'methane_water  1  1'
        fileout = 'meth_wat_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 5 #Initial time
        #t0disp = 6 # Initial time for display
        #t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Methane, $mmol/m^3$"
        text = 'maxval'
        #levmin = 0.
        #levmax = 5000.
        nlevplot = 20 # Number of countour levels
        #stepticks = 60
        pathsave = path
elif int(nplot) == 5: # Oxygen
        basename = 'oxygen_water  1  1'
        fileout = 'oxyg_wat_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 5 #Initial time
        #t0disp = 6 # Initial time for display
        #t1disp = 10.3 # Final time for display
        #nvardisp = nlevs_water
        title = "Oxygen, mg/l"
        #levmin = 0.
        #levmax = 10.
        nlevplot = 20 # Number of countour levels
        pathsave = path
elif int(nplot) == 6: # Carbon dioxide
        basename = 'co2_water  1  1'
        fileout = 'co2_wat_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 5 #Initial time
        #t0disp = 6 # Initial time for display
        #t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Carbon dioxide, $mmol/m^3$"
        #levmin = 0.
        #levmax = 20.
        nlevplot = 20 # Number of countour levels
        #stepticks = 2
        pathsave = path
elif int(nplot) == 7: # Methane bubble flux
        basename = 'fbblflx_ch4  1  1'
        fileout = 'methbubbleflx_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 5 #Initial time
        t0disp = 5 # Initial time for display
        t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Methane bubble flux, $mol/(m^2*s)$"
        pathsave = path
elif int(nplot) == 8: # Richardson number
        basename = 'Ri  1  1'
        fileout = 'Ri_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
        levmax = 1.
        #nvardisp = nlevs_water
        title = "Richardson number, $n/d$"
        pathsave = path
elif int(nplot) == 9: # Rp number
        basename = 'Rp  1  1'
        fileout = 'Rp_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
	levmin = 1.
        levmax = 100.
        #nvardisp = nlevs_water
        title = "R_p, $n/d$"
        pathsave = path
elif int(nplot) == 10: # Rp number
        basename = 'Rpdens  1  1'
        fileout = 'Rpdens_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
        #levmax = 1.
        #nvardisp = nlevs_water
        title = "$R_{p,\rho}, ~n/d$"
        pathsave = path
elif int(nplot) == 11: # x-speed component
        basename = 'u  1  1'
        fileout = 'u_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
        #levmax = 1.
	levmin = 0. #0.01
        #nvardisp = nlevs_water
        title = "x-speed component, $m/s$"
        pathsave = path
elif int(nplot) == 12: # y-speed component
        basename = 'v  1  1'
        fileout = 'v_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        #time0 = 5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
	#levmin = -0.01
        #levmax = 0.01
        #nvardisp = nlevs_water
        title = "y-speed component, $m/s$"
        pathsave = path
elif int(nplot) == 13: # Heat conductance
        basename = 'lamw  1  1'
        fileout = 'lamw_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 0. #5 #Initial time
        time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
        #nvardisp = nlevs_water
        title = "Heat conductance, $m^2/s$"
        pathsave = path
elif int(nplot) == 14: # epsilon, TKE dissipation rate
        basename = 'epsilon  1  1'
        fileout = 'epsilon_cont'
        nheader = 6 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
	time_unit = 1. #Time units in input data, expressed in days
        ctime = 5
        cstart = 7
        time0 = 0. #5 #Initial time
        #time0 = month_doy(7,ny)/mdays + 17./31.
        #t0disp = 4 # Initial time for display
        #t1disp = 10 # Final time for display
        #nvardisp = nlevs_water
        title = "TKE dissipation rate, $m^2/s^3$"
        pathsave = path
        logscalebar = True
elif int(nplot) == 15: # Temperature measured
        basename = 'KuivaTemp_201409_2' #'LDS'
        fileout = 'tempobs_cont'
        nheader = 0 #Number of lines in the header
        labelx = 'Time, days'
        labely = 'Depth, m'
        timescale = 24. # 365/12*24 - number of hours in a month
	time_unit = 1./24. #Time units in input data, expressed in days
        #ctime = 7
        #cstart = 9
	ctime = 2
        cstart = 4
        #time0 = 5 #Initial time
        time0 = 0.25 #month_doy(7,ny)/mdays + 17./31.
        #t0disp = 5 # Initial time for display
        #t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Temperature measured, $^\circ$C"       
        levmin = 8.
        levmax = 18.
        nlevplot = 20 # Number of countour levels
        pathsave = path
        path = pathobs
	delim = ','
elif int(nplot) == 16: # Oxygen measured
        basename = 'o2man2013cont'
        fileout = 'o2obsman_cont'
        nheader = 1 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
        ctime = 1
        cstart = 3
        time0 = 1 #Initial time
        t0disp = 5.9 # Initial time for display
        t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Oxygen measured, mg/l"
        levmin = 0.
        levmax = 10.
        nlevplot = 20 # Number of countour levels
        pathsave = path
        path = pathobs
elif int(nplot) == 17: # Methane measured
        basename = 'ch4man2013cont'
        fileout = 'ch4obsman_cont'
        nheader = 1 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
        ctime = 3
        cstart = 5
        time0 = 1 #Initial time
        t0disp = 6 # Initial time for display
        t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Methane measured, $\mu g/l$"
        levmin = 0.
        levmax = 600.
        text = 'maxval'
        nlevplot = 20 # Number of countour levels
	stepticks = 60
        pathsave = path
        path = pathobs
elif int(nplot) == 18: # Carbon dioxide measured
        basename = 'co2man2013cont'
        fileout = 'co2obsman_cont'
        nheader = 1 #Number of lines in the header
        labelx = 'Time, months'
        labely = 'Depth, m'
        timescale = 365/12 # 365/12*24 - number of hours in a month
        ctime = 3
        cstart = 5
        time0 = 1 #Initial time
        t0disp = 6 # Initial time for display
        t1disp = 11 # Final time for display
        #nvardisp = nlevs_water
        title = "Carbon dioxode measured, mg/l"
        levmin = 0.
        levmax = 20.
	stepticks = 2
        nlevplot = 20 # Number of countour levels
        pathsave = path
        path = pathobs



# Changing default matplotlib settings
plt.rc('xtick',labelsize=16)
plt.rc('ytick',labelsize=16)
plt.rc('lines',linewidth=2)

#nfiles = 10 #number of files
#time = [0]*nfiles
##Files names suffixes (timestamps)
#time[0] = '2003010100'
#time[1] = '2003010200'
#time[2] = '2003010300'
#time[3] = '2003010400'
#time[4] = '2003010500'
#time[5] = '2003010600'
#time[6] = '2003010700'
#time[7] = '2003010800'
#time[8] = '2003010900'
#time[9] = '2003011000'

#for i in range(nfiles):
fid = open(path + basename + '.dat','r')

# Counting number of levels
readheader(fid,nheader)
line = fid.readline()
if 'delim' in locals():
	line = line.split(delim)
else:
	line = line.split()
print(line)
nvardisp = int((len(line) - cstart)/2) + 1
fid.seek(0)
readheader(fid,nheader)

# Initializing arrays
ncol = nm.zeros(nvardisp,dtype="int32")
ncol[0] = cstart # The column where data starts
for i in range(nvardisp-1):
        ncol[i+1] = ncol[i] + 2
depths = nm.zeros(nvardisp)

#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 values is '+str(nraws))

fid.seek(0) #rewinding file
readheader(fid,nheader)

time = nm.zeros(nraws)
vals = nm.zeros((nraws,nvardisp)) #[[0]*nvardisp]*nraws

# Reading array of depths and values to be plotted
for j in range(nraws) :
        line = fid.readline()
	if 'delim' in locals():
        	line = line.split(delim)
	else:
        	line = line.split()
        time[j] = float(line[ctime-1])
        for i in range(nvardisp):
                vals[j,i] = float(line[ncol[i]-1])
                depths[i] = float(line[ncol[i]-2])
        
# Time scaling and shifting
if time_label_date:
        # Constructing time labels in date-time format
        time[:] = time[:]*time_unit + time0_
        time_ = []
        for i in range(len(time)):
                time_.append(dts.num2date(time[i]))
else:
        time[:] = time[:]/timescale + time0

# Calculating the time offset
if 't0disp' in locals():
        it0disp = timeoffset(time,t0disp)
else:
        it0disp = 0
if 't1disp' in locals():
        it1disp = timeoffset(time,t1disp)
else:
        it1disp = nraws-1

# Constructing contour levels
if 'levmax' in locals():
        levmax = levmax 
else :
        levmax = nm.nanmax(nm.nanmax(vals[it0disp:it1disp,:])) #nanmax, nanmin ignore NaNs
if 'levmin' in locals():
        levmin = levmin
else :
        levmin = nm.nanmin(nm.nanmin(vals[it0disp:it1disp,:]))

if (not logscalebar):
	#Linear scale
	levels = nm.zeros(nlevplot+1)
	step   = (levmax - levmin)/nlevplot
	levels[0] = levmin
	for i in range(nlevplot):
        	levels[i+1] = levels[i] + step
else:
	#Logarithmic scale
	maxbar = int(nm.log10(levmax) // 1 + 1)
	minbar = int(nm.log10(levmin) // 1 )
	nlevplot = 2*(maxbar - minbar)
	levels = nm.zeros(nlevplot+1)
	for i in range(nlevplot+1):
		levels[i] = 10.**(0.5*i+minbar)


print(levels)
# Plotting contour plot
if time_label_date:
        fig, ax = plt.subplots(1)
	if (logscalebar):
        	cont = ax.contourf(time_[:],depths,nm.transpose(vals[:,:]),levels,locator=ticker.LogLocator(),cmap=plt.cm.jet) 
	else:
        	cont = ax.contourf(time_[:],depths,nm.transpose(vals[:,:]),levels,cmap=plt.cm.jet) 
        plt.gca().xaxis.set_major_locator(days)
        plt.gca().xaxis.set_major_formatter(dateFmt)
        #fig.xaxis.set_minor_locator(months)
        fig.autofmt_xdate()
else:
        fig, ax = plt.subplots(1)
	if (logscalebar):
        	cont = ax.contourf(time[it0disp:it1disp],depths,nm.transpose(vals[it0disp:it1disp,:]),\
		levels,locator=ticker.LogLocator(),cmap=plt.cm.jet)
	else:
        	cont = ax.contourf(time[it0disp:it1disp],depths,nm.transpose(vals[it0disp:it1disp,:]),levels,cmap=plt.cm.jet)

plt.gca().invert_yaxis()
#print(depths)
#print(vals)

#Setting tick labels at x and y axes
#plt.xticks(nm.arange(mt.floor(time[it0disp]), mt.ceil(time[it1disp]), 1.0))
#plt.yticks(nm.arange(mt.floor(depths[0]), mt.ceil(depths[len(depths)-1]), 1.0))

#Setting labels
plt.xlabel(labelx,fontsize=20)
plt.ylabel(labely,fontsize=20)

# Drawing colourbar
if (not logscalebar):
	if not ('stepticks' in locals()):
	        stepticks = (levmax - levmin)/nlevplot
	        #norder = nm.floor(nm.log10(diff))
		#stepticks = round(diff/10**norder)*10**norder
	cbticks = [levels[0]]
	for i in range(nlevplot):
		# Adding only integer ticks
		nn = (levels[i+1] - levmin)/stepticks
       	 	#norder = nm.floor(nm.log10(nn))
       	 	x = nm.floor(nm.log10(abs(levels[i+1])))#Order of magnitude
		if abs(round(nn) - nn) < 10.**(x-3):
       	        	 #print(levels[i+1])
			cbticks.append(levels[i+1])
	plt.colorbar(cont,ticks=cbticks)
else:
	plt.colorbar(cont)

plt.title(title)


if 'text' in locals():
        if text == 'maxval':
                maxval = nm.nanmax(vals[it0disp:it1disp,:])
                plt.figtext(0.4, 0.8, text + ' = ' + \
		"%8.2f" % (maxval), \
                backgroundcolor='w',size=20,fontdict=None)

plt.savefig(pathsave + fileout + '_' + exp + '.png')
plt.show()
