# -*- coding: utf-8 -*-
# Program to read NEMO ocean model output temperature field file (netCDF3),
# Requirements:
# 1)convert it to netcdf4 and write one output variable per file.
# 2)Output files should also pass CF checker and CMIP conventions (and possibly OGC).
# 3) Mars attributes should be clearly identified 
# 4) time should be handled (forecast time, assimilation window etc)
# Use an approach whereby various attributes/additional variables can be added to output NetCDF files by calls to functions
#
# K Marsh 23/9/15
#
# Added option to handle multiple input files, input directory and output directory  
#
# K Marsh 2/6/16
#
# V2 is a rewritten version to streamline the code, as there is now no requirement to split the fields by level
# K Marsh 7/6/16
#
#*****************************************************************
from netCDF4 import Dataset
from netCDF4 import num2date
import datetime
import MySQLdb
import getopt
import sys
import glob


def add_ensemble(fid, ens_num, ens_len):
    fid.createDimension('ensemble', ens_len)
    fid.createVariable('realization', 'i', 'ensemble')
    fid.variables['realization'].setncattr('standard_name', 'realization')
    fid.variables['realization'].setncattr('long_name', "Number of the simulation in the ensemble")
    fid.variables['realization'][:] = ens_num
    return    

def add_mars_atts(fid, att_dict, var):
    st_date = fid.variables['atime'][:]
    date_str, time_str = str(num2date(st_date, fid.variables['atime'].units)).split()
    for item_att in att_dict.keys():
        if item_att == "mars_date" :
            print 'set date', date_str
            fid.variables[var].setncattr(item_att, date_str)
            continue
        if item_att == "mars_time":
            fid.variables[var].setncattr(item_att, time_str)
            continue
        if item_att == "mars_level":
            fid.variables[var].setncattr(item_att, str(fid.variables['deptht'][:]))
            continue        
        fid.variables[var].setncattr(item_att, att_dict[item_att])
    return

def modify_time(fid, var):
    lDebugFlag = True 
    fid.renameVariable('time_counter', 'time')
    fid.renameVariable('time_counter_bnds', 'time_bnds')
    print dir(fid)
    fid.renameDimension('time_counter','time')
    if lDebugFlag: 
        print fid.variables['time_bnds'][:]
    fid.variables['time'].setncattr("bounds", "time_bnds")
    data_var = fid.variables[var]
    for att in data_var.ncattrs():
        val = data_var.getncattr(att)
        print "att,val",att,val
        if "coordinates" in att and "time_counter" in val:
            val_new = val.replace("time_counter", "time")
            data_var.setncattr("coordinates", val_new)
            print "set time_counter to time"
    fid.createVariable('atime', 'd',)
    fid.variables['atime'].setncattr("standard_name", "forecast_reference_time")
    fid.variables['atime'].setncattr("units", fid.variables['time'].getncattr('units'))
    if lDebugFlag:
        print "** **",fid.variables['time_bnds'][0][0],fid.variables['time_bnds'][0][1]
    fid.variables['atime'][:]=fid.variables['time_bnds'][0][0]
    if lDebugFlag:
        print"***** fid atime ", fid.variables['atime'][:],fid.variables['atime'].shape
    return

def add_assimilation_window(fid, var):
    fid.createVariable('assimilation_window_start', 'd',)
    fid.variables['assimilation_window_start'].setncattr("long_name", "assimilation window start")
    fid.variables['assimilation_window_start'].setncattr("units", fid.variables['time'].getncattr('units'))
    fid.variables['assimilation_window_start'][:] = fid.variables['time_bnds'][0, 0]
    fid.createVariable('assimilation_window_end', 'd',)
    fid.variables['assimilation_window_end'].setncattr("long_name", "assimilation window end")
    fid.variables['assimilation_window_end'].setncattr("units", fid.variables['time'].getncattr('units'))
    fid.variables['assimilation_window_end'][:] = fid.variables['time_bnds'][0, 1]
    return

# Python script to perform a simple lookup on the ecmwf param database for a
# given paramID and return the corresponding CF Standard Name. Databases are
# opened as read-only
#


def db_lookup(db_name,host_name,user_name,pass_name,param_num):

    try:
        db = MySQLdb.connect(host = host_name, # your host, usually localhost
                     user = user_name, # your username
                      passwd = pass_name, # your password
                      db = db_name) # name of the data base

    except MySQLdb.Error, e:
  
        print "Error %d: %s" % (e.args[0],e.args[1])
        sys.exit(1)
#
# create a Cursor object to execute all the queries required
#
    cur = db.cursor() 
#
# List parameter from tables in the database
#

    print "Getting CF name for GRIB1 parameter number ", param_num
    paramID = str(param_num)
    query = "SELECT name, grib1_ecmwf FROM cf WHERE grib1_ecmwf = " + paramID
    cur.execute(query)
    desc = cur.description
    print "%s %3s" % (desc[0][0], desc[1][0])
#
# print all the first cell of all the rows
#
    for row in cur.fetchall() :
        print "CF Standard Name: %s  GRIB1 Code: %s \n" % (row[0],row[1])
    db.close()
    return row[0]

def read_config_file(lDebugFlag,cfg_fname):
    print "start!"
    try:
        cfg_file=open(cfg_fname,"r")
    except:
        print "File to process does not exit; Exiting"
        sys.exit()    
    

    lines=cfg_file.readlines()
    m_dict={}
    a_dict={}
    v_dict={}
    fn_list=[]
    o_dir=""
    for line in lines:
        if line.startswith("#") or len(line.strip()) == 0 : # comment/empty lines: 
            continue  
        elif line.startswith("\'"): # "Env: " line
            option, value = line.split("=")
            o=option.strip()
            v=value.strip()
            os.environ[o] ="$"+o+":"+v            
        else:
            if lDebugFlag:
                print line
            option, value = line.split("=")
            o=option.strip()
            v=value.strip()
            if o == "db_name":
                d_name = v
            elif o == "host_name":
                h_name = v
            elif o == "user_name":
                 u_name = v
            elif o == "pass_name":
                p_name =v
            elif o == "mars_type":
                 m_dict["mars_type"] = v                 
            elif o == "mars_class":
                m_dict["mars_class"] = v
            elif o == "mars_stream":
                 m_dict["mars_stream"] = v
            elif o == "mars_expver":
                m_dict["mars_expver"] =v
            elif o == "mars_date":
                 m_dict["mars_date"] = v                  
            elif o == "mars_time":
                m_dict["mars_time"] = v
            elif o == "mars_step":
                 m_dict["mars_step"] = v
            elif o == "mars_number":
                m_dict["mars_number"] =v
            elif o == "mars_levelist":
                 m_dict["mars_levelist"] = v   
            elif o == "mars_levtype":
                m_dict["mars_levtype"] = v
            elif o == "mars_param":
                 m_dict["mars_param"] = v
            elif o == "fnames":
                fn_list.append(v)
            elif o == "input_dir":
                fn_list = glob.glob(v+"\/*.nc")
            elif o == "output_dir": 
                o_dir = v
            elif o == "global_attributes_Conventions":
                 a_dict["Conventions"] = v              
            elif o == "global_attributes_title":
                 at_dict["title"] = v      
            elif o == "global_attributes_institution":
                 a_dict["institution"] = v      
            elif o == "global_attributes_references":
                 a_dict["references"] = v      
            elif o == "global_attributes_source":
                 a_dict["source"] = v  
            elif o == "global_attributes_history":
                 a_dict["history"] = v      
            elif o == "global_attributes_comment":
                 a_dict["comment"] = v                                                      
            elif o == "nemo_var_name":
                 v_dict[v.split(":")[0]]= [v.split(":")[1],v.split(":")[2],v.split(":")[3]]   
            else:
                print "Warning: Unmatched option"
    cfg_file.close()
    if len(fn_list) == 0:
        print "No files to process: exiting"
        sys.exit()
    if o_dir == "":
        print "No output directory specified: exiting"
        sys.exit()      
    if not o_dir.endswith("\/"):
        o_dir =o_dir + "/"
          
    if len(m_dict) != 11:
        print "Missing Mars elements: exiting"
        sys.exit()
    if lDebugFlag:
        print "MARS dictionary is :", m_dict  
    return d_name, h_name,u_name,p_name,m_dict,fn_list,o_dir,a_dict,v_dict




def main():
    """
     main part of the program; take an output NEMO file and split by variable
     
     1 parse config file
     2 get file list
     3 process files in list
     4 get required variables in list
     5 get required axes
     6 get required attributes
     7 get coordinates
     8 get list of timesteps
     9 write out by timesteps
     10 amend variable attributes
     11 amend global attributes
     12 modify hours (NEMO offset)
     13 write out file
     
     
    """
#
# Read configuration file; also build dictionary to hold MARS-specific 
# placeholder metadata 
#
    lDebugFlag = False
    options, remainder = getopt.getopt(sys.argv[1:],'c:')
    if len(options) ==0:
        print "Option required; Exiting"
        sys.exit()
    for opt, arg in options:
        print opt,arg
        if opt in ('-c'):
           cf_cfg_fname = arg
           print "Config file is ",cf_cfg_fname
        else:
            print "Config file required; Exiting"
            sys.exit()
    db_name, host_name,user_name,pass_name,mars_dict,fnames, op_dir,at_d,var_dict = read_config_file(lDebugFlag, cf_cfg_fname)
    print  db_name, host_name,user_name,pass_name,mars_dict,fnames, op_dir,at_d,var_dict
#
# input filename list for testing; input and output directories are taken from the config file
#
    file_count = 0
    for fname in fnames:
        file_count = file_count + 1
        print fname
        fname_out = op_dir+"a"+(fname.split("/")[-1]).split(".")[0]
        dsin = Dataset(fname, "r")
        print fname_out
#
#copy existing file global attributes to output files
#
        nc_global_attrs = dsin.ncattrs()
        print nc_global_attrs
#        
# list of files to process -only those specified in config file with CF standard name       
#        
        var_list = []
        for v_name, varin in dsin.variables.iteritems():
            if v_name in var_dict.keys():
                var_list.append(str(v_name)) 
        print "Variables in file : ",var_list
#
# This is the required list of variables needed in all output 'split' files:
#
        process_list = [item for item in var_list if item in var_dict.keys()]
        print "variables to process : ",process_list
        print len(var_list)," variables in the file "
        print len(var_dict.keys()), " Variables specified in config file"
        print len(process_list)," Variables to process " 
        nvar = 0
        icount = 0
#
# loop over required variables 
#
        for p_name in process_list:
           bounds_list=[]
           icount = icount + 1
#
# get variable from input file
#
           var = dsin.variables[p_name]
           print var
           print dir(var)
           print "processing ", icount, " ", p_name
           print "writing to "+ fname_out + "_reformat_time_cat_" + p_name + ".nc"
           dsout = Dataset(fname_out + "_reformat_time_cat_" + p_name + ".nc", "w", format = "NETCDF3_CLASSIC")
#
# create dimensions; use list of dimension names (strings), and take dimensions from input files (dsin.dimensions)
#
           dim_list = var.dimensions
           str_dim_list = [str(i) for i in dim_list]           
           print "dim list is ",dim_list
           print str_dim_list
           for dname, the_dim in dsin.dimensions.iteritems():
              print dname, len(the_dim),the_dim,dir(the_dim)
              if dname in dim_list or dname == "tbnds": 
                  print dname, "match"
                  dsout.createDimension(dname, len(the_dim) if not the_dim.isunlimited() else None)
#
# Add coordinate variables and auxiliary coordinate variables; 
# loop through variables in file to get required variables
#
           for v_name, varin in dsin.variables.iteritems():
               if v_name in str_dim_list or v_name in var.coordinates or v_name == "time_counter_bnds":
                   print "add coordinate variables", v_name
                   outVar = dsout.createVariable(v_name, varin.datatype, varin.dimensions)
#                  
# Copy coordinate variable attributes
#
                   outVar.setncatts({k: varin.getncattr(k) for k in varin.ncattrs()})
                   for j in varin.ncattrs():
                       if j == "bounds":
                           print v_name, " has bounds"
                           print dir(j)
                   outVar[:] = varin[:]                 
                   print "set coordinate values and attributes for ", v_name          

           print "adding variable ", var._name
           print var._name, var.datatype, var.dimensions
#
# Option to rename netcdf variables to 'data' rather than the NEMO name
#
           set_name_to_data = 1
           if set_name_to_data == 0:
               variable_name=var._name
           if set_name_to_data == 1:
               variable_name = "data"
               NEMO__netcdf_variable_name = var._name
           outVar1 = dsout.createVariable(variable_name, var.datatype, var.dimensions)
           outVar1[:] = var[:]    
#
# Copy variable attributes
#
           for k in var.ncattrs():
#
#ignore title attributes for variables as this contravenes CF
#ignore _FillValue attributes for variables as this is not set correctly
#                    
              if k == "_FillValue" or k == "title": 
                  continue
              if k == "standard_name":
                  outVar1.setncattr("NEMO_standard_name", var.getncattr(k))
                  continue
              if k == "coordinates":
                  outVar1.setncattr(k, "atime " + var.getncattr(k))
                  continue
              outVar1.setncatts({k: var.getncattr(k) })
#
# Todo: set variable fillvalue to original rather than default
#               

#
# if forecast fields, modify time axis accordingly
#
           modify_time(dsout, variable_name)               
           print "output file contents ",dir(dsout.variables)
           print dsout.variables.keys()
#
# Add global attributes from original file
# Delete existing history
#
           print len(nc_global_attrs)
           if len(nc_global_attrs) > 0:
               for ncattr in nc_global_attrs:
                  if ncattr == "history":
                      continue
                  dsout.setncattr(ncattr, nc_global_attrs[ncattr])
#
# Set CF standard names 
#
           cf_name=var_dict[p_name][0]
           print "CF Name is ",cf_name
           grib_param_id = var_dict[p_name][1]
           grib_param_id_str = str(grib_param_id)
           mars_dict['mars_param'] = grib_param_id_str
           add_mars_atts(dsout, mars_dict, "data")
           outVar1.setncattr("standard_name",cf_name)
           dsout.close()
           dsin.close()
           break

"""
#
# This is the required list of variables needed in all output 'split' files:
#
        req_list = ['nav_lon', 'nav_lat','deptht','time_counter', 'time_counter_bnds']
        process_list = [item for item in var_list if item not in req_list]
        nvar = 0
        icount = 0
#
# loop over required variables which are not required in every file
#
        for p_name in process_list:
            icount = icount + 1
#
# loop over variable items in the netcdf file; use string rather than unicode 
#
            for v_name, varin in dsin.variables.iteritems():
                v_name_str = str(v_name)
#
# if a surface filed, skip for now as need to rework the dimensions
#
                if v_name_str[0:2] == "so": 
                    print "skip surface field ", v_name_str
                    continue
#
# if matches a process variable name, open a new file for output
#
                if v_name_str == p_name:
                    print "processing ", icount, " ", v_name_str, type(v_name_str)
                    print "writing to "+ fname_out + "_reformat_time_cat_" + v_name_str + ".nc"
                    dsout = Dataset(fname_out + "_reformat_time_cat_" + v_name_str + ".nc", "w", format = "NETCDF3_CLASSIC")
#
# add global attributes from original file
#
                    if len(nc_global_attrs) > 0:
                        for ncattr in nc_global_attrs:
                            dsout.setncattr(ncattr, nc_global_attrs[ncattr])

#
# write out the dimensions
#
                    for dname, the_dim in dsin.dimensions.iteritems():
                        dsout.createDimension(str(dname), len(the_dim)) 
                    add_ensemble(dsout, file_count, 1)
#
# iterate again to write out the required co-ordinate variables
#
                    for v_name_req, varin_req in dsin.variables.iteritems():        
                        if str(v_name_req) in req_list:
                            if "_FillValue" in varin_req.ncattrs() :                      
                                f_val = getattr(varin_req, "_FillValue")
                                outVar = dsout.createVariable(str(v_name_req), \
			        varin_req.dtype, varin_req.dimensions, fill_value = f_val)
                            else:
                                outVar = dsout.createVariable(str(v_name_req), \
			        varin_req.dtype, varin_req.dimensions)
#
# Writing out required co-ordinate Variable: 
#
                            outVar[:] = varin_req[:]
#
# Have to set variable attributes one at a time; _FillValue was not set correctly; 
# title is not a CF variable attribute
#
                            for ncattr in varin_req.ncattrs():
                                if ncattr == "_FillValue" or ncattr == "title": 
                                    continue
#
# TODO add CF standard name attribute based on NEMO-> CF mapping 
#
                                if ncattr == "standard_name":
                                    outVar.setncattr("NEMO_standard_name", varin_req.getncattr(ncattr))
                                    continue                           
#
# Setting attributes for required co-ordinate variable attribute: 
#
                                outVar.setncattr(ncattr, varin_req.getncattr(ncattr))
#
# write out the 'to process' variable  - v_name_str is the original variable name; 
# _FillValue has to be handled carefully and set in the createVariable call.
#
                    if "_FillValue" in varin.ncattrs():
                        varin1 = ['ensemble',]
                        for x in varin.dimensions:
                            varin1.append(x)
                        f_val = getattr(varin, "_FillValue")
                        outVar = dsout.createVariable("data", varin.dtype, varin1, fill_value = f_val)
                    else:
                        outVar = dsout.createVariable("data", varin.dtype, varin1)
                    print "Writing output variable", v_name_str,varin1, " as \"data\"" 
                    outVar[:] = varin[:] 
#
# Have to set variable attributes one at a time
#
                    outVar.setncattr("NEMO_variable_name", v_name_str)
                    for ncattr in varin.ncattrs():

                    
#
# Using dummy ID -hardcoded in script as not all cf names have grib mapping
# add cf standard_name placeholder
#
                    cf_name=var_dict[v_name][0]
                    print "CF Name is ",cf_name
                    grib_param_id = var_dict[v_name][1]
                    grib_param_id_str = str(grib_param_id)
                    mars_dict['mars_param'] = grib_param_id_str
#                    cf_name=db_lookup(db_name,host_name,user_name,pass_name,grib_param_id_str)
                    add_mars_atts(dsout, mars_dict, "data")
                    outVar.setncattr("standard_name",cf_name)
                    print dsout.dimensions
                    print "closing output file", fname_out
                    dsout.close()
                    f_out=fname_out + "_reformat_time_cat_" + v_name_str + ".nc"
                    print f_out
#                    dsout1 = Dataset(f_out,"a")
#                    print dsout1.data_model
#                    print dsout1.dimensions
#                    dsout1.renameDimension(u'time_counter',u'time')
#                    dsout1.close()
                    nvar = nvar + 1
                    print "Variables processed/total: ",nvar,len(process_list)
  
#
# have processed this 'process' variable, so exit loop and do the next
#
                    if lDebugFlag:
                        break
#
# for testing, just need to output the first two variables from the input file
#
            if lDebugFlag: 
                break
        print "Variables Processed/Total: ", nvar,len(process_list)

"""
if __name__ == "__main__":
    main()




