# -*- 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
#
#*****************************************************************
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 main():
    """
     main part of the program; take an output NEMO file and split by variable
    """
#
# 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()
    try:
        cfg_file=open(cf_cfg_fname,"r")
    except:
        print "File to process does not exit; Exiting"
        sys.exit()
    lines=cfg_file.readlines()
    mars_dict={}
    at_d={}
    var_dict={}
    fnames=[]
    op_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":
                db_name = v
            elif o == "host_name":
                host_name = v
            elif o == "user_name":
                 user_name = v
            elif o == "pass_name":
                pass_name =v
            elif o == "mars_type":
                 mars_dict["mars_type"] = v                 
            elif o == "mars_class":
                mars_dict["mars_class"] = v
            elif o == "mars_stream":
                 mars_dict["mars_stream"] = v
            elif o == "mars_expver":
                mars_dict["mars_expver"] =v
            elif o == "mars_date":
                 mars_dict["mars_date"] = v                  
            elif o == "mars_time":
                mars_dict["mars_time"] = v
            elif o == "mars_step":
                 mars_dict["mars_step"] = v
            elif o == "mars_number":
                mars_dict["mars_number"] =v
            elif o == "mars_levelist":
                 mars_dict["mars_levelist"] = v   
            elif o == "mars_levtype":
                mars_dict["mars_levtype"] = v
            elif o == "mars_param":
                 mars_dict["mars_param"] = v
            elif o == "fnames":
                fnames.append(v)
            elif o == "input_dir"
                fnames = glob.glob(v+"\/*.nc")
            elif o == "output_dir" 
                op_dir = v
            elif o == "global_attributes_Conventions":
                 at_d["Conventions"] = v              
            elif o == "global_attributes_title":
                 at_d["title"] = v      
            elif o == "global_attributes_institution":
                 at_d["institution"] = v      
            elif o == "global_attributes_references":
                 at_d["references"] = v      
            elif o == "global_attributes_source":
                 at_d["source"] = v  
            elif o == "global_attributes_history":
                 at_d["history"] = v      
            elif o == "global_attributes_comment":
                 at_d["comment"] = v                                                      
            elif o == "nemo_var_name":
                 var_dict[v.split(":")[0]]= [v.split(":")[1],v.split(":")[2],v.split(":")[3]]   
            else:
                print "Warning: Unmatched option"
    cfg_file.close()
    if len(fnames) == 0:
        print "No files to process: exiting"
        sys.exit()
    if op_dir == "":
        print "No output directory specified: exiting"
        sys.exit()      
    if op_dir.endswith("\/"):
        continue
    else:
        op_dir =op_dir + "/"
          
    if len(mars_dict) != 11:
        print "Missing Mars elements: exiting"
        sys.exit()
    if lDebugFlag:
        print "MARS dictionary is :", mars_dict
#
# get current date for 'History' global attribute
#
    now = datetime.datetime.now()
    date_now = ("%s" % now)
    date_now_iso = ("%s" % now.isoformat())
    date_now_history = ("%s-%s-%s %s:%s" % (now.year, now.month, now.day, now.hour, now.minute ))

#
#input filename list for testing
#
# files are output to /scratch/ma/makm/nemo_cfg/split_data_cf/
#
    file_count = 0
    for fname in fnames:
        file_count = file_count + 1
        fname_out = op_dir+"a"+(fname.split("/")[-1]).split(".")[0]
        dsin = Dataset(fname, "r")
#
#copy existing file global attributes to output files
#
        nc_global_attrs = dsin.ncattrs()
#        
# 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 to process: ",var_list
#
# 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])
#
# Add global metadata
# Delete existing history
#
                    if "history" in at_d:
                       del at_d["history"]

                    for key in at_d:
                        dsout.setncattr(key, at_d[key])
#
# 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():
#
#ignore title attributes for variables as this contravenes CF
#ignore _FillValue attributes for variables as this is not set correctly
#                    
                        if ncattr == "_FillValue" or ncattr == "title": 
                            continue
                        if ncattr == "standard_name":
                            outVar.setncattr("NEMO_standard_name", varin.getncattr(ncattr))
                            continue
                        if ncattr == "coordinates":
                            outVar.setncattr(ncattr, "atime " + varin.getncattr(ncattr))
                            continue
                        outVar.setncattr(ncattr, varin.getncattr(ncattr))
#
# if forecast fields, modify time axis accordingly
#
                    modify_time(dsout, "data")
                    
#
# 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()



