#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as dates
import datetime
from matplotlib.dates import date2num
import pandas
import netCDF4
from pylab import savefig
import sys, getopt

def main(argv):
  ncfile = ''
  outputdir = ''
  errmsg='Use: '+str(sys.argv[0])+' -i <ncfile> -o <outputdir>'
  try:
     opts, args = getopt.getopt(argv,"h:i:o:")
  except getopt.GetoptError:
     print errmsg
     sys.exit(2)
  for opt, arg in opts:
     if opt == '-h':
        print errmsg
        sys.exit()
     elif opt in ("-i"):
        ncfile = arg
     elif opt in ("-o"):
        outputdir = arg
  if len(ncfile) == 0 or len(outputdir) == 0:
    print 'Please specify an input file and an output directory.'
    print errmsg
    sys.exit()
  print 'Input file is ', ncfile
  print 'Output directory is ', outputdir

  # MAIN PROGRAM
  # -----------------------------------------------------------------
  
  # DATA FROM THE MAST ----------------------------------------------
  
  # Read the files using the (very convenient) Pandas reader
  data = pandas.read_csv('vent09+_GABLS4.dat',sep=';', na_values="99.9", skiprows=13)
  # Add the index axis
  data.index = pandas.to_datetime(data['Date'], format='%Y-%m-%d %H:%M:%S')
  data = data[45:115]
  date_obs = data.index
  #levels_obs = [3.5, 10.9, 18.3, 25.6, 33., 42.2] # temp lev
  levels_obs = [4.15, 9.67, 18.87, 26.23, 33.59, 41.99] # wind lev
  
  
  # GCM RESULTS -----------------------------------------------------
  
  decomp=ncfile.rsplit('/',5)
  version=decomp[len(decomp)-3]
  nc = netCDF4.Dataset(ncfile)
  vitu = nc.variables['vitu'][:]
  vitv = nc.variables['vitv'][:]
  windv = (vitu**2. + vitv**2.)**0.5
  time = nc.variables['time_counter'][:]
  lati_id = 0
  longi_id = 0
  alti_id = range(4)
  date = netCDF4.num2date(time[:], units = 'seconds since 2009-12-11 08:00:00')
  alti_var = np.zeros(4)
  flabel = ['']*4
  if str(nc.variables).find("geop") > -1 and \
     str(nc.variables).find("phis") > -1:
    geop = nc.variables['geop'][:]
    phis = nc.variables['phis'][:]
    for ilev in range(4):
      alti_var[ilev] = np.mean(geop[:,alti_id[ilev],lati_id,longi_id] - \
                       phis[:,lati_id,longi_id])/9.8
      flabel[ilev] = \
        "LMDz (z="+str("%.1f" % alti_var[ilev])+"m)"
  else:
    flabel[0] = "LMDz (z=6m approx.)"
    flabel[1] = "LMDz (z=20m approx.)"
    flabel[2] = "LMDz (z=35m approx.)"
    flabel[3] = "LMDz (z=53m approx.)"
  
  # -----------------------------------------------------------------
  
  # DATA VISUALIZATION ----------------------------------------------
  # We plot the figure
  
  fig = plt.figure()
  ax = fig.gca()
  
  linegcm=plt.plot(date,windv[:,alti_id[0],lati_id,longi_id],'k-o',linewidth=2,label=flabel[0]) 
  # GCM
  lineobs=plt.plot(date_obs,data['vm1'],'k--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[0])+'m') # DATA monthly mean
  lineobs=plt.plot(date_obs,data['vm2'],'k--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[1])+'m') # DATA monthly mean
  linegcm=plt.plot(date,windv[:,alti_id[1],lati_id,longi_id],'b-o',linewidth=2,label=flabel[1]) 
  # GCM
  lineobs=plt.plot(date_obs,data['vm3'],'b--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[2])+'m') # DATA monthly mean
  lineobs=plt.plot(date_obs,data['vm4'],'b--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[3])+'m') # DATA monthly mean
  linegcm=plt.plot(date,windv[:,alti_id[2],lati_id,longi_id],'r-o',linewidth=2,label=flabel[2]) 
  # GCM
  lineobs=plt.plot(date_obs,data['vm5'],'r--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[4])+'m') # DATA monthly mean
  lineobs=plt.plot(date_obs,data['vm6'],'r--', \
    linewidth=2,label='OBS '+str("%.1f" % levels_obs[5])+'m') # DATA monthly mean
  
  #ax.xaxis.set_major_locator(dates.HourLocator())
  #ax.xaxis.set_major_locator(dates.DayLocator())
  #ax.xaxis.set_major_locator(dates.MonthLocator())
  ax.xaxis.set_major_formatter(dates.DateFormatter('%Y-%m-%d %H:%M:%S'))
  # Maximum content = '%Y-%b-%d %H:%M:%S'
  #ax.xaxis.set_major_formatter(dates.DateFormatter('%d %b %Y'))
  #ax.xaxis.set_major_formatter(dates.DateFormatter('%b %Y'))
  ax.set_title("Wind speed, GABLS4, "+version)
  #ax.set_title("Top (red) and bottom (blue) - LMDz")
  #ax.set_title("Data (blue) and LMDz (red) - Bottom")
  #ax.set_ylim([-0.2,80.])
  ax.set_ylabel("Wind speed (m/s)")
  #ax.set_ylabel("Relative humidity over ice")
  #ax.set_ylabel("Partial pressure of water vapor (Pa)")
  #ax.set_yscale('log')
  #ax.set_xlabel("Date")
  plt.grid()
  #plt.plot(codedate, uservar2, 'r-')
  #plt.legend(('I (surface)', 'I (sommet)'), loc=(1.03, 0.8))
  #plt.legend(('pvap (bottom)', 'pvap (top)'), loc=(1.03, 0.8))
  plt.xticks(rotation=30,ha='right') #plt.xticks(rotation='vertical')
  plt.subplots_adjust(bottom=0.2)
  plt.subplots_adjust(right=0.8) # keep room for legend e.g. 0.8
  plt.subplots_adjust(left=0.2) # keep room for legend e.g. 0.8
  
  #plt.show()
  savefig('atlas_wind_'+version+'.pdf', bbox_inches='tight')
  # -----------------------------------------------------------------
  # -----------------------------------------------------------------
  # -----------------------------------------------------------------

# MAIN PROGRAM
# Just calling the main function
if __name__ == "__main__":
   main(sys.argv[1:])

#------------------------------------------------------------------