# Python script to transfomr ASCII LIDAR outputs to netCDF
## g.e. # LIDAR_ASCII_netCDF.py -f /d4/lflmd/etudes/REMEMBER/WRFmeas/tests/WL_HyMeX/gir.d00.01DAR

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

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

fillValue = 1.e20

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

parser = OptionParser()
parser.add_option("-f", "--LIDAR_file", dest="lfile",
                  help="LIDAR ASCII text file to use", metavar="FILE")

(opts, args) = parser.parse_args()

lidarvn = ['z', 'p', 'u', 'v', 'w', 't_pot', 'qv', 'qc', 'qr', 'qs', 'qh', 'qi',     \
  'qg', 'dens','cldfra']
lidarvln = ['height', 'pressure','x-wind direction', 'y-wind direction',             \
  'z-wind direction', 'potential temperature', 'water vapor mixing ratio',           \
  'cloud mixing ratio', 'rain mixing ratio', 'snow mixing ratio',                    \
  'hail mixing ratio', 'ice mixing ratio', 'graupel mixing ratio', 'air density',    \
  'cloud fraction']
lidarvu = ['m', 'hPa', 'ms-1', 'ms-1', 'ms-1', 'K', 'kgkg-1', 'kgkg-1', 'kgkg-1',    \
  'kgkg-1', 'kgkg-1', 'kgkg-1', 'kgkg-1', 'kg m-3','1']

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

ofile = 'lidar.nc'
Nlidarvariables = len(lidarvn)

if not os.path.isfile(opts.lfile):
    print errormsg
    print '  LIDAR ASCII text file "' + opts.lfile + '" does not exist !!'
    print errormsg
    quit()

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', 'LIDAR_ASCII_netCFD.py')
ncvar.set_attribute(objofile, 'version', '1.0')

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

lidarvals = {}

iline=0
dimz = 0
Searchdimz = True
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')
        ncvar.set_attribute(newvar, 'name',values[0])
        ncvar.set_attribute(newvar, 'acronym',values[3])
        ncvar.set_attribute(newvar, 'real_lon',                                      \
          np.float(values[6].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'real_lat',                                      \
          np.float(values[5].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'x_grid_point',                                  \
          int(values[8].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'y_grid_point',                                  \
          int(values[9].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'model_lon',                                     \
          np.float(values[12].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'model_lat',                                     \
          np.float(values[11].replace(',','').replace('(','').replace(')','')) )
        ncvar.set_attribute(newvar, 'model_height',                                  \
          np.float(values[13].replace(',','').replace('(','').replace(')','')) )
        simstarttime = values[18]
    else:
        if values[0] == 'new_time':
            time_step.append(np.float(values[2]))
            psfc.append(np.float(values[6]))
            rainc.append(np.float(values[7]))
            rainnc.append(np.float(values[8]))
            drydens.append(np.float(values[9]))
            if iline == 1: 
                Searchdimz = True
            else:
                Searchdimz = False
        else:
            if not values[0] == 'k':
                lidarvals[itz] = values
                if Searchdimz:
                    dimz = dimz + 1
                itz = itz + 1
    iline = iline + 1

dimt = len(time_step)

print '  Found:',dimt,'time steps, over:',dimz,'vertical levels'
objlfile.close()

objofile.createDimension('z',dimz)

time_stepv = np.zeros((dimt), dtype=np.float)
psfcv = np.zeros((dimt), dtype=np.float)
raincv = np.zeros((dimt), dtype=np.float)
rainncv = np.zeros((dimt), dtype=np.float)
drydensv = np.zeros((dimt), dtype=np.float)

lidarvaluesv = np.zeros( (dimt,dimz,Nlidarvariables), dtype= np.float)

itz = 0
for it in range(dimt):
    time_stepv[it] = np.float(time_step[it])
    psfcv[it] = np.float(psfc[it])
    raincv[it] = np.float(rainc[it])
    rainncv[it] = np.float(rainnc[it])
    drydensv[it] = np.float(drydens[it])

    for iz in range(dimz):
        for iv in range(Nlidarvariables):
            lidarvaluesv[it,iz,iv] = np.float(lidarvals[itz][iv+1])

        itz = itz + 1
# Surface variables
newvar = objofile.createVariable('time','f4',('time',))
newvar[:] = time_stepv
newattr = ncvar.basicvardef(newvar, 'time', 'time', 'hours since ' +                 \
  simstarttime.replace('_',' '))

newvar = objofile.createVariable('psfc','f4',('time',))
newvar[:] = psfcv
newattr = ncvar.basicvardef(newvar, 'psfc', 'surface pressure', 'hPa')

newvar = objofile.createVariable('rainc','f4',('time',))
newvar[:] = raincv
newattr = ncvar.basicvardef(newvar, 'rainc',                                         \
  'accumulated precipitation from cumulus scheme', 'mm')

newvar = objofile.createVariable('rainnc','f4',('time',))
newvar[:] = rainncv
newattr = ncvar.basicvardef(newvar, 'rainnc',                                        \
  'accumulated precipitation not from cumulus scheme', 'mm')

newvar = objofile.createVariable('drydens','f4',('time',))
newvar[:] = drydensv
newattr = ncvar.basicvardef(newvar, 'drydens', 'total dry air column pressure', 'hPa')

# Lidar variables
for iv in range(Nlidarvariables):
    newvar = objofile.createVariable(lidarvn[iv], 'f4', ('time','z'))
    newvar[:] = lidarvaluesv[:,:,iv]
    newattr = ncvar.basicvardef(newvar, lidarvn[iv], lidarvln[iv], lidarvu[iv] )

objofile.sync()
objofile.close()

print 'Successfull generation of LIDAR netCDF file "' + ofile + '" !!!!!'