#! /usr/bin/env python
#
# ===========================================================
#
# Alexander R Davies
#
# Ocean Exploration, Remote Sensing, and Biogeography Lab
# School of Marine Science and Policy
# College of Earth, Ocean, and Environment
# University of Delaware
# ardavies@udel.edu
#
# Latest Update: 10/09/13
#
# NOTES: 	1) Reads Processed Data, Linearly Interpolates 
#				each for each profile to a known grid
#
#			2) Calculates vertical velocities for density
#				and chlorphyll. 
#
#			3) Plots the vertical velocities over original
#				data
#
# ===========================================================
# Import Modules
# ===========================================================
#
import numpy as np
import csv
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import glob
import os
from matplotlib.ticker import MaxNLocator
from matplotlib import rc, rcParams
from numpy.random import uniform, seed
from matplotlib.mlab import griddata
from numpy import arange,array,ones#,random,linalg
from pylab import plot,show
from scipy import interpolate
from matplotlib.colors import LogNorm
from matplotlib.backends.backend_pdf import PdfPages
#
# ===========================================================
# Function to read csv files
# ===========================================================
#
def csvread(filename):
	with open(filename, 'rb') as f:
	    reader = csv.reader(f, delimiter=',')
	    nrow = 0;
	    for row in reader:
	    	nrow = nrow +1
	    	ncol = len(row)
	array = np.zeros([ncol,nrow]) 
	with open(filename, 'rb') as f:
	    reader = csv.reader(f, delimiter=',')
	    counter = 0
	    for row in reader:
	    	for jj in range(0,ncol):
	    		array[jj,counter] = row[jj]
	    	counter = counter + 1
	return array, nrow, ncol   
#
# ===========================================================
# Function to Find the nearest value
# ===========================================================
#
def find_nearest(array,value):
    idx = (np.abs(array-value)).argmin()
    return array[idx], idx
#
# ===========================================================
# Sort Data and Descriptive Arrays
# ===========================================================
#
os.chdir('C:/Documents and Settings/Alex/Documents/Research/Delaware/MUSTACHE/Jan01_Jun04_2013/data/type')
fullnames = glob.glob('Full*')
gradnames = glob.glob('Grad*')
infonames = glob.glob('Text*')
arraylen =  len(fullnames)
floatdate = np.zeros(arraylen)
lat = [None]*arraylen
lon = [None]*arraylen
day = np.zeros(arraylen)
month=[None]*arraylen
year=[None]*arraylen
daystr=[None]*arraylen
#
# ===========================================================
# Build the Arrays as days past 1/1/13
# ===========================================================
#
for d in range(0,int(arraylen)):
	f = open(infonames[d], 'r')
	while 1:
		line = f.readline()
		if line[:4] == "Date":
			month[d] = str(line[5:7])
			day[d] = int(line[8:10])
			daystr[d] = str(line[8:10])
			year[d] = str(line[11:15])
			break
	f = open(infonames[d], 'r')
	while 1:
		line = f.readline()
		if line[:3] == "Lat":
			lat[d] = str(line[5:12])
			break
	f = open(infonames[d], 'r')
	while 1:
		line = f.readline()
		if line[:3] == "Lon":
			lon[d] = str(line[5:12])
			break
	
	if (year[d] == '2013'):
		yr = 0
	elif (year[d] == '2014'):
		yr = 365
	elif (year[d] == '2015'):
		yr = 365*2
	if(month[d] == '01'):
		mon = 0
	elif(month[d] == '02'):
		mon = 31
	elif(month[d] == '03'):
		mon = 31 + 28
	elif(month[d] == '04'):
		mon = (31 + 28 + 31)
	elif(month[d] == '05'):
		mon = (31 + 28 + 31 + 30)
	elif(month[d] == '06'):
		mon = (31 + 28 + 31 + 30 + 31)
	elif(month[d] == '07'):
		mon = (31 + 28 + 31 + 30 + 31 + 30)
	elif(month[d] == '08'):
		mon = (31 + 28 + 31 + 30 + 31 + 30 + 31)
	elif(month[d] == '09'):
		mon = (31 + 28 + 31 + 30 + 31 + 30 + 31 + 31)
	elif(month[d] == '10'):
		mon = (31 + 28 + 31 + 30 + 31 + 30 + 31 + 31 + 30)
	elif(month[d] == '11'):
		mon = (31 + 28 + 31 + 30 + 31 + 30 + 31 + 31 + 30 + 31)
	elif(month[d] == '12'):
		mon = (31 + 28 + 31 + 30 + 31 + 30 + 31 + 31 + 30 + 31 + 30)
	else:
		print 'invalid month'

	floatdate[d] = yr + mon + day[d]
#
# ===========================================================
# Read the Data and Interpolate for Gridding
# ===========================================================
#
Ygrid = np.linspace(-1800,-10,1791*2)
interplength = len(Ygrid)
GridData = np.zeros([interplength,7,arraylen])
ChlyDumby = np.zeros([interplength,arraylen])
for i in range(0,arraylen):
	dataout, rows, cols = csvread(fullnames[i])
	#
	TemInterp = interpolate.interp1d(dataout[0,:], dataout[1,:],kind='linear')
	DenInterp = interpolate.interp1d(dataout[0,:], dataout[2,:],kind='linear')
	SalInterp = interpolate.interp1d(dataout[0,:], dataout[3,:],kind='linear')
	PreInterp = interpolate.interp1d(dataout[0,:], dataout[4,:],kind='linear')
	OxyInterp = interpolate.interp1d(dataout[0,:], dataout[7,:],kind='linear')
	ChlInterp = interpolate.interp1d(dataout[0,:], dataout[11,:],kind='linear')
	BskInterp = interpolate.interp1d(dataout[0,:], dataout[12,:],kind='linear')
	CDMInterp = interpolate.interp1d(dataout[0,:], dataout[13,:],kind='linear')
	#
	for j in range(0,interplength):
		GridData[j,0,i] = TemInterp(Ygrid[j])
		GridData[j,1,i] = DenInterp(Ygrid[j])
		GridData[j,2,i] = SalInterp(Ygrid[j])
		GridData[j,3,i] = OxyInterp(Ygrid[j])			 
		ChlyDumby[j,i] = ChlInterp(Ygrid[j])	
		GridData[j,5,i] = BskInterp(Ygrid[j])	
		GridData[j,6,i] = CDMInterp(Ygrid[j])	
		if (ChlInterp(Ygrid[j]) < 0):
			GridData[j,4,i] = GridData[j-1,4,i]
		else: 
			GridData[j,4,i] = ChlyDumby[j,i]	
#
# ===========================================================
# Find Gridded Change Velocities Based on Nearest Values
# ===========================================================
#
GridWData = np.zeros([interplength,7,arraylen-1])
GridWDepths = np.zeros([interplength,7,arraylen-1])
GridWDates = np.zeros([interplength,7,arraylen-1])
for k in range(0,7):
	for i in range (0,arraylen-1):
		for j in range(0,interplength):
			value, index = find_nearest(GridData[:,k,i+1], GridData[j,k,i])
			startdepth = Ygrid[j]
			newdepth = Ygrid[index]
			GridWDepths[j,k,i] = (newdepth + startdepth)/2
			dz = newdepth-startdepth
			GridWData[j,k,i] = dz/172800
			GridWDates[j,k,i] = (floatdate[i] + floatdate[i+1])/2
#
# ===========================================================
# Find Avg Gridded Change Velocities Based on Nearest Values
# ===========================================================
#
GridWDataAvg = np.zeros([interplength,7,arraylen-3])
GridWDepthsAvg = np.zeros([interplength,7,arraylen-3])
GridWDatesAvg = np.zeros([interplength,7,arraylen-3])
for k in range(0,7):
	for i in range (0,arraylen-3):
		for j in range(0,interplength):
			GridWDataAvg[j,k,i] = (GridWData[j,k,i] + GridWData[j,k,i+1] + GridWData[j,k,i+2])/3
			GridWDepthsAvg[j,k,i] = (GridWDepths[j,k,i] + GridWDepths[j,k,i+1] + GridWDepths[j,k,i+2])/3
			GridWDatesAvg[j,k,i] = (GridWDates[j,k,i] + GridWDates[j,k,i+1] + GridWDates[j,k,i+2])/3
#
# ===========================================================
# Find Depth of Isochly Contours
# ===========================================================
#
isobiolines = np.linspace(.1,.40,20)
numbiolines = len(isobiolines)
Biodepths = np.zeros([numbiolines,arraylen])
for i in range(0,arraylen):
	#print i
	for j in range(0,numbiolines):
		#print j
		count = 0
		while (GridData[count,4,i] < isobiolines[j]):
				count = count + 1 
		Biodepths[j,i] = Ygrid[count]
#
# ===========================================================
# Find Isochly Contour Change Velocities
# ===========================================================
#
IsoBioWData = np.zeros([numbiolines,arraylen-1])
IsoBioWDepths = np.zeros([numbiolines,arraylen-1])
IsoBioWDates = np.zeros([numbiolines,arraylen-1])
IsoBioFake = np.zeros([numbiolines,arraylen-1])
for i in range(0,arraylen-1):
	for k in range(0,numbiolines):
		j = i
		IsoBioWData[k,i] = (Biodepths[k,j+1] - Biodepths[k,j])/172800
		IsoBioFake[k,i] = 0
		IsoBioWDepths[k,i] = (Biodepths[k,j+1] + Biodepths[k,j])/2
		IsoBioWDates[k,i] = (floatdate[j+1] + floatdate[j])/2
#
# ===========================================================
# Find Avg Isoline Change Velocities
# ===========================================================
#
IsoBioWDataAvg = np.zeros([numbiolines,arraylen-3])
IsoBioWDepthsAvg = np.zeros([numbiolines,arraylen-3])
IsoBioWDatesAvg = np.zeros([numbiolines,arraylen-3])
IsoBioFakeAvg = np.zeros([numbiolines,arraylen-3])
for i in range (0,arraylen-3):
	for j in range(0,numbiolines):
			IsoBioWDataAvg[j,i] = (IsoBioWData[j,i] + IsoBioWData[j,i+1] + IsoBioWData[j,i+2])/3
			IsoBioWDepthsAvg[j,i] = (IsoBioWDepths[j,i] + IsoBioWDepths[j,i+1] + IsoBioWDepths[j,i+2])/3
			IsoBioWDatesAvg[j,i] = (IsoBioWDates[j,i] + IsoBioWDates[j,i+1] + IsoBioWDates[j,i+2])/3
			IsoBioFakeAvg[j,i] = (IsoBioFake[j,i] + IsoBioFake[j,i+1] + IsoBioFake[j,i+2])/3

# ===========================================================
# Find Sink Vel
# ===========================================================
#
SinkWData = np.zeros([numbiolines,arraylen-1])
for i in range (0,arraylen-1):
	for k in range(0,numbiolines):
		value, index = find_nearest(GridWDepths[:,1,i], IsoBioWDepths[k,i])
		SinkWData[k,i] = IsoBioWData[k,i] - GridWData[index,1,i]
#
# ===========================================================
# Find Avg Sink Vel
# ===========================================================
#
SinkWDataAvg = np.zeros([numbiolines,arraylen-3])
for i in range (0,arraylen-3):
	for k in range(0,numbiolines):
		value, index = find_nearest(GridWDepthsAvg[:,1,i], IsoBioWDepthsAvg[k,i])
		SinkWDataAvg[k,i] = IsoBioWDataAvg[k,i] - GridWDataAvg[index,1,i]
#		
# ===========================================================
# Chly Contour Map w/ Chly Isoline Ws
# ===========================================================
#
figdirpath ='C:/Documents and Settings/Alex/Documents/Research/Delaware/MUSTACHE/Jan01_Jun04_2013/plots/contours_with_Ws/'
#
# Plot Initialization
contlevels= np.logspace(-2.5,.65,100)
conticks = [0.01,0.25,0.50, 1.00, 1.50, 2.00, 3.00]
figtit = 'Chlorophyll Contour (micro gram/liter); Isochly W Quiver'
savetit = "Chlorophyll_Contour_ChlorophyllW_Quivers"
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
#
# Plotting correct contour data
contourdat = np.zeros([interplength,arraylen])
for ci in range(0,interplength):
	for cj in range(0,arraylen):
		contourdat[ci,cj] = GridData[ci,4,cj]
#
# Contour Plot
cs = plt.contourf(floatdate,Ygrid,contourdat, levels= contlevels, norm = LogNorm())
cbar = plt.colorbar(cs, spacing='proportional', norm = LogNorm())
cbar.set_ticks(conticks)
cbar.set_ticklabels(conticks)
#
# Quiver Plot
from pylab import *
Q = plt.quiver(IsoBioWDates[1::10], IsoBioWDepths[1::10], IsoBioFake[1::10], IsoBioWData[1::10])
plt.quiverkey(Q, 0.22, 0.95, 0.001, r'$5 \times 10^{-3} m/s$', labelpos='W',
               fontproperties={'weight': 'bold'})
#
# Labels
ax.set_xlabel("Days since 01/01/2013")
ax.set_ylabel("Depth (m)")
ax.set_ylim([-900,0])
fig.suptitle(figtit) 
plt.savefig(figdirpath + savetit + ".png")
#
# ===========================================================
# Chly Contour Map w/ Avg Chly Isoline Ws
# ===========================================================
#
figdirpath ='C:/Documents and Settings/Alex/Documents/Research/Delaware/MUSTACHE/Jan01_Jun04_2013/plots/contours_with_Ws/'
#
# Plot Initialization
contlevels= np.logspace(-2.5,.65,100)
conticks = [0.01,0.25,0.50, 1.00, 1.50, 2.00, 3.00]
figtit = 'Chlorophyll Contour (micro gram/liter); 3 Profile Avg Isochly W Quiver'
savetit = "Chlorophyll_Contour_ChlorophyllW_Quivers_3Avg"
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
#
# Plotting correct contour data
contourdat = np.zeros([interplength,arraylen])
for ci in range(0,interplength):
	for cj in range(0,arraylen):
		contourdat[ci,cj] = GridData[ci,4,cj]
#
# Contour Plot
cs = plt.contourf(floatdate,Ygrid,contourdat, levels= contlevels, norm = LogNorm())
cbar = plt.colorbar(cs, spacing='proportional', norm = LogNorm())
cbar.set_ticks(conticks)
cbar.set_ticklabels(conticks)
#
# Quiver Plot
from pylab import *
Q = plt.quiver(IsoBioWDatesAvg[1::10], IsoBioWDepthsAvg[1::10], IsoBioFakeAvg[1::10], IsoBioWDataAvg[1::10])
plt.quiverkey(Q, 0.22, 0.95, 0.001, r'$5 \times 10^{-3} m/s$', labelpos='W',
               fontproperties={'weight': 'bold'})
#
# Labels
ax.set_xlabel("Days since 01/01/2013")
ax.set_ylabel("Depth (m)")
ax.set_ylim([-900,0])
fig.suptitle(figtit) 
plt.savefig(figdirpath + savetit + ".png")

os.chdir('C:/Documents and Settings/Alex/Documents/Research/Delaware/MUSTACHE/Jan01_Jun04_2013/code/interpolate_profiles')