#!/usr/bin/env python
# Last modified:  Time-stamp: <2008-12-17 15:37:00 haines>
"""ims_sodar_plot_month"""

# plot each month

import os, sys, glob, re
import datetime, time, dateutil, dateutil.tz
import pycdf
import numpy

sys.path.append('/home/haines/nccoos/raw2proc')
del(sys)

os.environ["MPLCONFIGDIR"]="/home/haines/.matplotlib/"

from pylab import figure, twinx, savefig, setp, getp, cm, colorbar
from matplotlib.dates import DayLocator, HourLocator, MinuteLocator, DateFormatter, date2num, num2date
import procutil

print 'ims_sodar_plot_month ...'

proc_dir = '/seacoos/data/nccoos/level1/ims/sodar/'
fns = glob.glob((os.path.join(proc_dir, '*.nc')))
fns.sort()

print len(fns)

for fn in fns[14:16]:
    m=re.search('\d{4}_\d{2}', fn)
    yyyy_mm = m.group()
    prev_month, this_month, next_month = procutil.find_months(yyyy_mm)
    
    # load data
    print ' ... ... read: ' + fn
    nc = pycdf.CDFMF((fn,))
    ncvars = nc.variables()
    # print ncvars
    es = nc.var('time')[:]
    units = nc.var('time').units
    dt = [procutil.es2dt(e) for e in es]
    # set timezone info to UTC (since data from level1 should be in UTC!!)
    dt = [e.replace(tzinfo=dateutil.tz.tzutc()) for e in dt]
    # return new datetime based on computer local
    dt_local = [e.astimezone(dateutil.tz.tzlocal()) for e in dt]
    dn = date2num(dt)
    z = nc.var('z')[:]
    # convert cm/s to m/s
    uu = nc.var('u')[:]/100.
    vv = nc.var('v')[:]/100.
    ww = nc.var('w')[:]/100.
    echo = nc.var('echo')[:]
    n1 = numpy.floor(nc.var('nois1')[:]/10.)
    n2 = numpy.floor(nc.var('nois2')[:]/10.)
    n3 = numpy.floor(nc.var('nois3')[:]/10.)
    nc.close()

    # retain original dt for addnan
    dto = dt

    # last dt in data for labels
    dt1 = dt[-1]
    dt2 = dt_local[-1]
    
    diff = abs(dt1 - dt2)
    if diff.days>0:
        last_dt_str = dt1.strftime("%H:%M %Z on %b %d, %Y") + ' (' + dt2.strftime("%H:%M %Z, %b %d") + ')'
    else:
        last_dt_str = dt1.strftime("%H:%M %Z") + ' (' + dt2.strftime("%H:%M %Z") + ')' \
                      + dt2.strftime(" on %b %d, %Y")

    fig = figure(figsize=(10, 10))
    fig.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, wspace=0.1, hspace=0.1)

#######################################
# Plot month
#######################################

    ax = fig.add_subplot(5,1,1)
    axs = [ax]

    # range for horizontal wind plots
    cmin, cmax = (-20., 20.)
    # print "%s : %g %g" % ('uv wind', cmin, cmax)
    # use masked array to hide NaN's on plot
    (dt, uu) = procutil.addnan(dto, uu)
    dn = date2num(dt)
    um = numpy.ma.masked_where(numpy.isnan(uu), uu)
    pc = ax.pcolor(dn, z, um.T, vmin=cmin, vmax=cmax)
    pc.set_label('True Eastward Wind (m s-1)')
    ax.text(0.025, 0.1, pc.get_label(), fontsize="small", transform=ax.transAxes)
    ax.set_ylabel('Height (m)')
    # ax.set_ylim(-26.,2.)

# setup colorbar axes instance.
    l,b,w,h = ax.get_position()
    cax = fig.add_axes([l, b+h+0.04, 0.25*w, 0.03])

    cb = colorbar(pc, cax=cax, orientation='horizontal') # draw colorbar
    cb.set_label('Wind Velocity (m s-1)')
    cb.ax.xaxis.set_label_position('top')
    cb.ax.set_xticks([0.1, 0.3, 0.5, 0.7, 0.9])
    xtl = numpy.round(numpy.linspace(cmin, cmax, 10), decimals=0)
    cb.ax.set_xticklabels([xtl[1], xtl[3], xtl[5], xtl[7], xtl[9]])
    #cb.ax.set_xticklabels([-16, -8, 0, 8, 16])

# ax.set_xlim(dt[0], dt[-1]) # first to last regardless of what  
    ax.set_xlim(date2num(this_month), date2num(next_month-datetime.timedelta(seconds=1))) 
    ax.xaxis.set_major_locator( DayLocator(range(2,32,2)) )
    ax.xaxis.set_minor_locator( HourLocator(range(0,25,12)) )
    ax.set_xticklabels([])

# this only moves the label not the tick labels
    ax.xaxis.set_label_position('top')
    ax.set_xlabel('IMS SODAR -- ' + yyyy_mm)

# right-hand side scale
    ax2 = twinx(ax)
    ax2.yaxis.tick_right()
# convert (lhs) meters to (rhs) feet
    feet = [procutil.meters2feet(val) for val in ax.get_ylim()]
    ax2.set_ylim(feet)
    ax2.set_ylabel('Height (ft)')

#######################################
#
    ax = fig.add_subplot(5,1,2)
    axs.append(ax)

# use masked array to hide NaN's on plot
    (dt, vv) = procutil.addnan(dto, vv)
    dn = date2num(dt)
    vm = numpy.ma.masked_where(numpy.isnan(vv), vv)
    pc = ax.pcolor(dn, z, vm.T, vmin=cmin, vmax=cmax)
    pc.set_label('True Northward Wind (m s-1)')
    ax.text(0.025, 0.1, pc.get_label(), fontsize="small", transform=ax.transAxes)
    ax.set_ylabel('Height (m)')
    # ax.set_ylim(-26.,2)

    # ax.set_xlim(date2num(dt[0]), date2num(dt[-1])) 
    ax.set_xlim(date2num(this_month), date2num(next_month-datetime.timedelta(seconds=1))) 
    ax.xaxis.set_major_locator( DayLocator(range(2,32,2)) )
    ax.xaxis.set_minor_locator( HourLocator(range(0,25,12)) )
    ax.xaxis.set_major_formatter( DateFormatter('%m/%d') )
    ax.set_xticklabels([])

    # right-hand side scale
    ax2 = twinx(ax)
    ax2.yaxis.tick_right()
    # convert (lhs) meters to (rhs) feet
    feet = [procutil.meters2feet(val) for val in ax.get_ylim()]
    ax2.set_ylim(feet)
    ax2.set_ylabel('Height (ft)')
    
    #######################################
    #
    ax = fig.add_subplot(5,1,3)
    axs.append(ax)
    
    # range for horizontal wind plots
    cmin, cmax = (-5., 5.)
    # print "%s : %g %g" % ('w wind', cmin, cmax)

    # use masked array to hide NaN's on plot
    (dt, ww) = procutil.addnan(dto, ww)
    dn = date2num(dt)
    wm = numpy.ma.masked_where(numpy.isnan(ww), ww)
    pc = ax.pcolor(dn, z, wm.T, vmin=cmin, vmax=cmax)
    # pc.set_label('Upward Wind (m s-1)')
    ax.text(0.025, 0.1, pc.get_label(), fontsize="small", transform=ax.transAxes)
    ax.set_ylabel('Height (m)')
    # ax.set_ylim(-26.,2)
    
    # setup colorbar axes instance.
    l,b,w,h = ax.get_position()
    cax = fig.add_axes([l+0.04, b+h-0.06, 0.25*w, 0.03])

    cb = colorbar(pc, cax=cax, orientation='horizontal') # draw colorbar
    cb.set_label('Upward Wind Velocity (m s-1)')
    cb.ax.xaxis.set_label_position('top')
    cb.ax.set_xticks([0.1, 0.3, 0.5, 0.7, 0.9])
    xtl = numpy.round(numpy.linspace(cmin, cmax, 10), decimals=0)
    cb.ax.set_xticklabels([xtl[1], xtl[3], xtl[5], xtl[7], xtl[9]])

    # ax.set_xlim(date2num(dt[0]), date2num(dt[-1])) 
    ax.set_xlim(date2num(this_month), date2num(next_month-datetime.timedelta(seconds=1))) 
    ax.xaxis.set_major_locator( DayLocator(range(2,32,2)) )
    ax.xaxis.set_minor_locator( HourLocator(range(0,25,12)) )
    ax.xaxis.set_major_formatter( DateFormatter('%m/%d') )
    ax.set_xticklabels([])
    
    # right-hand side scale
    ax2 = twinx(ax)
    ax2.yaxis.tick_right()
    # convert (lhs) meters to (rhs) feet
    feet = [procutil.meters2feet(val) for val in ax.get_ylim()]
    ax2.set_ylim(feet)
    ax2.set_ylabel('Height (ft)')
    
    #######################################
    #
    ax = fig.add_subplot(5,1,4)
    axs.append(ax)
    
    cmin, cmax = (0., 1000.)
    # print "%s : %g %g" % ('echo', cmin, cmax)

    # use masked array to hide NaN's on plot
    (dt, echo) = procutil.addnan(dto, echo)
    dn = date2num(dt)
    em = numpy.ma.masked_where(numpy.isnan(echo), echo)
    pc = ax.pcolor(dn, z, em.T, vmin=cmin, vmax=cmax)
    # pc.set_label('Echo Strength ()')
    ax.text(0.025, 0.1, pc.get_label(), fontsize="small", transform=ax.transAxes)
    ax.set_ylabel('Height (m)')
    # ax.set_ylim(-26.,2)

    # setup colorbar axes instance.
    l,b,w,h = ax.get_position()
    cax = fig.add_axes([l+0.04, b+h-0.06, 0.25*w, 0.03])

    cb = colorbar(pc, cax=cax, orientation='horizontal') # draw colorbar
    cb.set_label('Echo Strength')
    cb.ax.xaxis.set_label_position('top')
    cb.ax.set_xticks([0.0, 1.0])
    xtl = numpy.round(numpy.linspace(cmin, cmax, 10), decimals=0)
    cb.ax.set_xticklabels([str(cmin), str(cmax)])
    
    # ax.set_xlim(date2num(dt[0]), date2num(dt[-1])) 
    ax.set_xlim(date2num(this_month), date2num(next_month-datetime.timedelta(seconds=1))) 
    ax.xaxis.set_major_locator( DayLocator(range(2,32,2)) )
    ax.xaxis.set_minor_locator( HourLocator(range(0,25,12)) )
    ax.xaxis.set_major_formatter( DateFormatter('%m/%d') )
    ax.set_xticklabels([])
    
    
    # right-hand side scale
    ax2 = twinx(ax)
    ax2.yaxis.tick_right()
    # convert (lhs) meters to (rhs) feet
    feet = [procutil.meters2feet(val) for val in ax.get_ylim()]
    ax2.set_ylim(feet)
    ax2.set_ylabel('Height (ft)')
    
    
    #######################################
    #
    ax = fig.add_subplot(5,1,5)
    axs.append(ax)
        
    # ax.plot returns a list of lines, so unpack tuple
    (dt, n1) = procutil(dto, n1)
    l1, = ax.plot_date(dt, n1, fmt='b-')
    l1.set_label('n1 each beam')
    (dt, n2) = procutil(dto, n2)
    l2, = ax.plot_date(dt, n2, fmt='c-')
    l2.set_label('n2')    
    (dt, n3) = procutil(dto, n3)
    l3, = ax.plot_date(dt, n3, fmt='g-')
    l3.set_label('n3')
    
    ax.set_ylabel('Ambient\n Noise')
    ax.set_ylim(20,120)
    ax.set_xlim(date2num(this_month), date2num(next_month-datetime.timedelta(seconds=1)))
    ax.xaxis.set_major_locator( DayLocator(range(2,32,2)) )
    ax.xaxis.set_minor_locator( HourLocator(range(0,25,12)) )
    ax.xaxis.set_major_formatter( DateFormatter('%m/%d') )

    ax.set_xlabel('IMS SODAR -- ' + yyyy_mm)
    
    # legend
    ls1 = l1.get_label()
    ls2 = l2.get_label()
    ls3 = l3.get_label()
    leg = ax.legend((l1,l2,l3), (ls1,ls2,ls3), loc='upper left')
    ltext  = leg.get_texts()  # all the text.Text instance in the legend
    llines = leg.get_lines()  # all the lines.Line2D instance in the legend
    frame  = leg.get_frame()  # the patch.Rectangle instance surrounding the legend
    frame.set_facecolor('0.80')      # set the frame face color to light gray
    frame.set_alpha(0.5)             # set alpha low to see through
    setp(ltext, fontsize='small')    # the legend text fontsize
    setp(llines, linewidth=1.5)      # the legend linewidth
    # leg.draw_frame(False)           # don't draw the legend frame

    # save figure
    
    ofn = '/home/haines/rayleigh/img/ims/ims_sodar1_'+yyyy_mm+'.png'
    print '... ... write: %s' % (ofn,)
    savefig(ofn)