# Python script to transfomr ASCII LIDAR outputs to netCDF
## g.e. # TS_ASCII_netCDF.py -f //media/ExtDiskD/bkup/ciclad/etudes/WL_HyMeX/iop15/wrf/run/control/stations_20121018000000-20121022000000/h0001.d01.TS -s 20121018000000

import os
from optparse import OptionParser
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import nc_var_tools as ncvar

main = 'TS_ASCII_netCDF.py'
errormsg='ERROR -- error -- ERROR -- error'
warnmsg='WARNING -- warning -- WARNING -- warning'

fillValue = 1.e20

def ts_header(ln):
    """ Function to get the values of the header of the *.TS files
      line=ASCII lines with the header of the TS file
      getting the line format from WRFV3.3 'EMCORE' in file 'share/wrf_timeseries.F'
    """
    fname = 'ts_header'

    fmts=['A26', 'I2', 'I3', 'A6', 'A2', 'F7.3', 'A1', 'F8.3', 'A3', 'I4', 'A1', 'I4',\
      'A3', 'F7.3', 'A1', 'F8.3', 'A2', 'F6.1', 'A7']

    headervalues = ncvar.values_fortran_fmt(ln,fmts)

    return headervalues

####### ###### ##### #### ### ## #

parser = OptionParser()
parser.add_option("-f", "--TS_file", dest="lfile",
  help="Time Series ASCII text file to use", metavar="FILE")
parser.add_option("-s", "--SimulationStartTime", dest="stime",
  help="Starting time of the simulation ([YYYY][MM][DD][HH][MI][SS] format)", metavar="DATE")

(opts, args) = parser.parse_args()


tsvn = ['t', 'q', 'u', 'v', 'psfc', 'glw', 'gsw', 'hfx', 'lh', 'tsk', 'tslb1', 'rainc', 'rainnc', 'clw']

tsvln = ['2 m Temperature', '2 m vapor mixing ratio', '10 m U wind (earth-relative)', '10 m V wind (earth-relative)', 'surface pressure', 'downward longwave radiation flux at the ground (downward is positive)', 'net shortwave radiation flux at the ground (downward is positive)', 'surface sensible heat flux (upward is positive)', 'surface latent heat flux (upward is positive)', 'skin temperature', 'top soil layer temperature', 'rainfall from a cumulus scheme', 'rainfall from an explicit scheme', 'total column-integrated water vapor and cloud variables']

tsvu = ['K', 'kg/kg', 'm/s', 'm/s', 'Pa', 'W/m2', 'W/m2', 'W/m2', 'W/m2', 'K', 'K', 'mm', 'mm', '1']

#######    #######
## MAIN
    #######

ofile = 'ts.nc'
Ntsvariables = len(tsvn)

if not os.path.isfile(opts.lfile):
    print errormsg
    print '  ' + main + ': Time-Series ASCII text file "' + opts.lfile +             \
      '" does not exist !!'
    print errormsg
    quit()

if opts.stime is None:
    print errormsg
    print '  ' + main + ': No initial date/time of the simulation is provided!'
    quit(-1)
else:
    stime = opts.stime
    refdate = stime[0:4] + '-' + stime[4:6] + '-' + stime[6:8] + ' ' + stime[8:10] + \
      ':' + stime[10:12] + ':' + stime[12:14]

objlfile = open(opts.lfile, 'r')

objofile = NetCDFFile(ofile, 'w')

# Creation of dimensions
##
objofile.createDimension('time',None)

ncvar.set_attribute(objofile, 'author', 'Lluis Fita Borrell')
ncvar.set_attribute(objofile, 'institution', 'Laboratoire Meteorologique Dynamique')
ncvar.set_attribute(objofile, 'university', 'University Pierre et Marie Curie')
ncvar.set_attribute(objofile, 'center', 'Centre national de la recherche scientifique')
ncvar.set_attribute(objofile, 'country', 'France')
ncvar.set_attribute(objofile, 'city', 'Paris')
ncvar.set_attribute(objofile, 'script', 'TS_ASCII_netCFD.py')
ncvar.set_attribute(objofile, 'version', '1.0')

time_step = []
psfc = []
rainc = []
rainnc = []
drydens = []

tsvals = {}

iline=0
itz = 0
for line in objlfile:
    values = ncvar.reduce_spaces(line)
#    print iline, values[0], dimz, Searchdimz
# Writting general information
    if iline == 0:
        newvar = objofile.createVariable('station','c')
        valueshead = ts_header(line)

        ncvar.set_attribute(newvar, 'name', ncvar.reduce_last_spaces(valueshead[0]))
        ncvar.set_attribute(newvar, 'acronym',valueshead[3].replace(' ',''))

        ncvar.set_attribute(newvar, 'real_lon', valueshead[5])
        ncvar.set_attribute(newvar, 'real_lat', valueshead[7])

        ncvar.set_attribute(newvar, 'x_grid_point', valueshead[9])
        ncvar.set_attribute(newvar, 'y_grid_point', valueshead[11])

        ncvar.set_attribute(newvar, 'model_lon', valueshead[13])
        ncvar.set_attribute(newvar, 'model_lat', valueshead[15])
        ncvar.set_attribute(newvar, 'model_height', valueshead[17])
        simstarttime = refdate
    else:
        tsvals[itz] = values
        time_step.append(np.float(values[1]))
        itz = itz + 1
    iline = iline + 1

dimt = len(time_step)

print '  Found:',dimt,'time steps'
objlfile.close()

time_stepv = np.zeros((dimt), dtype=np.float)
tsvaluesv = np.zeros( (dimt,Ntsvariables), dtype= np.float)

pracc = np.zeros((dimt), dtype=np.float)

itz = 0
for it in range(dimt):
    time_stepv[it] = np.float(time_step[it])

    for iv in range(Ntsvariables):
        tsvaluesv[it,iv] = np.float(tsvals[itz][iv+5])

    pracc[it] = np.float(tsvals[it][16]) + np.float(tsvals[it][17])
    itz = itz + 1

# time
newvar = objofile.createVariable('time','f8',('time'))
newvar[:] = time_stepv*3600.
newattr = ncvar.basicvardef(newvar, 'time', 'time', 'seconds since ' +               \
  simstarttime.replace('_',' '))
ncvar.set_attribute(newvar, 'calendar', 'standard')

# time-series variables
for iv in range(Ntsvariables):
    newvar = objofile.createVariable(tsvn[iv], 'f4', ('time'))
    newvar[:] = tsvaluesv[:,iv]
    newattr = ncvar.basicvardef(newvar, tsvn[iv], tsvln[iv], tsvu[iv] )

# Extra vars

# pr
varvals = np.zeros((dimt), dtype=np.float)
varvals[1:dimt] = pracc[1:dimt] - pracc[0:dimt-1]

newvar = objofile.createVariable('pr', 'f4', ('time'))
newvar[:] = varvals
newattr = ncvar.basicvardef(newvar, 'pr', 'precipitation', 'mm' )

objofile.sync()
objofile.close()

print 'Successfull generation of Time-Series netCDF file "' + ofile + '" !!!!!'