#!/usr/bin/env python
#
# -*- 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.
# 3) Mars attributes should be added clearly identified 
# 4) Modifications to time should be handled 
# 
# K Marsh 20/7/16
#
#*****************************************************************
# /convert_nemo_to_1_variable_all_levels_and_timesteps.py -c convert_nemo_to_1_variable_all_levels_and_timesteps.cfg
#
#
from netCDF4 import Dataset,netcdftime,num2date
import glob
import sys
import getopt
import datetime
import sys

def process_config_file(cfg_file_in):
    lDebugFlag = False
    lines = cfg_file_in.readlines()
    mars_dict = {}
    at_dict = {}
    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_dict["Conventions"] = v              
            elif o == "global_attributes_title":
                 at_dict["title"] = v      
            elif o == "global_attributes_institution":
                 at_dict["institution"] = v      
            elif o == "global_attributes_references":
                 at_dict["references"] = v      
            elif o == "global_attributes_source":
                 at_dict["source"] = v  
            elif o == "global_attributes_history":
                 at_dict["history"] = v      
            elif o == "global_attributes_comment":
                 at_dict["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", o
    cfg_file_in.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 not op_dir.endswith("\/"):
        op_dir =op_dir + "/"
          
    if len(mars_dict) != 11:
        print "Missing Mars elements: exiting"
        sys.exit()
    if lDebugFlag:
        print "MARS dictionary is :", mars_dict
    return fnames, at_dict, mars_dict, var_dict, op_dir


def get_original_dimensions_and_global_attributes(file_in):
#
#get dimensions
#
    dim_list = file_in.dimensions
    dim_dict={}
#
#Copy dimensions from whole file
#
    for dname, the_dim in file_in.dimensions.iteritems():
        dim_dict[dname] = len(the_dim)
#
# get global attributes 
#
    nc_attrs = file_in.ncattrs()
    print nc_attrs
    print "NetCDF Global Attributes:"
    if len(nc_attrs) == 0:
        print "None"
    for nc_attr in nc_attrs:
         print nc_attr           
    return dim_dict, nc_attrs
#   
# change existing variable attributes -form a dictionary to hold the ones to be used
#

def modify_variable_attributes(var_in):
    lDebugFlag1 = True
    lDebugFlag1 = False
    var_atts_out = {}
    var_atts_to_skip=["title",]
    for ncattr in var_in.ncattrs():    
        if "standard_name" in ncattr:
            var_atts_out["NEMO_standard_name"] = var_in.getncattr(ncattr)
            continue
        if ncattr == "_FillValue":
            f_val = getattr(var_in, "_FillValue")
            var_atts_out[ncattr] = f_val
            continue
        if ncattr == "title": 
            continue
        if lDebugFlag1:
            print ncattr       
            print var_in.getncattr(ncattr)
        var_atts_out[ncattr] = var_in.getncattr(ncattr)    
    if lDebugFlag1:
        print "variable attributes" , var_atts_out
    return var_atts_out
        
def add_mars_atts(fid, att_dict, var):

    for item_att in att_dict.keys():
        print item_att
#
#  need to skip mars_param attribute for cases when no mapping?
#             
        print item_att, type(item_att)
        print att_dict[item_att], type(att_dict[item_att]) 
        if item_att in ["mars_levelist"]:
           fid.variables[var].setncattr(item_att, float(att_dict[item_att]))
        elif item_att in ["mars_number","mars_time","mars_expver","mars_date"]:
            fid.variables[var].setncattr(item_att, int(att_dict[item_att]))
        else:
            fid.variables[var].setncattr(item_att, att_dict[item_att])
    return        

#
# need to adjust the reference time to that specified in filename
#

def adjust_time(fname_in,time_units_orig):
   lDebugFlag = False
   filename_time = (fname_in.split("/")[-1]).split("_")
   time1 = filename_time[2][-2:]
   time2 = filename_time[3][-2:]
   print time1,time2
   if time1 == time2 :
       time_units_parts = []
       time_units_parts = time_units_orig.split()
       time_units_parts[-1] = time1 + ":00:00" 
       time_units_new = " ".join(time_units_parts)
       if lDebugFlag:
           print time_units_new
   else: 
       print "times in filenames do not match: exit"
       sys.exit(1)
   return time_units_new
        
def main():
    lDebugFlag = False
    print "to run: convert_nemo_to_1_variable_all_levels_one_timestep.py -c <config file>"
    print "e.g. convert_nemo_to_1_variable_all_levels_one_timestep.py -c convert_nemo_to_1_variable_all_levels_one_timesteps.cfg"
    print "Parsing input config file"
    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()
    
#  Initialise variables to hold file information

    files = []
    at_cfg_dict = {}
    mars_cfg_dict = {}
    var_cfg_dict = {}
    start_date = ""
    start_time = ""
    
    files, at_cfg_dict, mars_cfg_dict, var_cfg_dict, output_dir = process_config_file(cfg_file)
 
    if lDebugFlag:
        print "at_cfg_dict", at_cfg_dict 
        print mars_cfg_dict 
        print var_cfg_dict
        print files
        print "files to process: "
        print files,
        print "\n"
    
#    files = ["/hugetmp/cera20c/split/2366_1d_1900022721_1900022821_grid_T_02.nc",] # dummy for testing
    
    if len(files) == 0:
       print "No files to process - exiting "
       sys.exit()

    for fname in files:
        fin=Dataset(fname,'r')
        file_dims, file_atts = get_original_dimensions_and_global_attributes(fin)
        print "Input file Dimensions: ",file_dims
        print "Input file global attributes: ", file_atts

# get list of variables

        fin_var_list = []
        for v in fin.variables:
            if lDebugFlag:
                print v
            fin_var_list.append(v)

        print fin_var_list
        var_count = 0
        for vname in fin_var_list:

# loop through variables

####            if vname == "votemper":  #temp var to select -use below in operations

            if vname in var_cfg_dict.keys():
                print "match for ", vname 
                var=fin.variables[vname]

# get variable attributes

                var_atts={k: var.getncattr(k) for k in var.ncattrs()}
                if lDebugFlag:
                    print "Variable attributes",
                    print var_atts
                    print "Variable Dimensions"
                    print var.dimensions, len(var.dimensions)
                for i in range(len(var.dimensions)):
                    if lDebugFlag:
                        print type(var.dimensions[i])
                        print i, var.dimensions[i]
                    if "time" in str(var.dimensions[i]):
                        time_index = i
                        time_length = var.shape[i]
                        if lDebugFlag:
                            print "Variable shape is ",var.shape
                            print "Variable time match at index:", i
                            print time_length
                            print "Variable coordinates"
                            print var.coordinates
                            print var.coordinates.split()[i]
                        if "time" in var.coordinates.split()[i]:
                            time_values = fin.variables[var.coordinates.split()[i]]
                            if lDebugFlag: 
                                print "time in coordinates at correct position"
                                print time_values[:]

# if timestep >1 loop through timesteps
# get timestep assuming time dimension is the first/second  dimension element

                        for time_step in range(time_length):
                            if time_index == 0:
                                var_time_step=var[time_step:time_step+1,...]

                            if time_index == 1:
                                var_time_step=var[:,time_step:time_step+1,...]            
                            time_step_value=fin.variables['time_counter'][time_step]        
                            if lDebugFlag: 
                                print "/nVariable: time step value", time_step_value
                                print "/nVariable: shape is ", var_time_step.shape

# get var coordinates, attributes 

                            if lDebugFlag: 
                                print "\noutput variable dimensions:"
                                print var.dimensions,var_time_step.shape
                            var_time_step_dimensions={}
                            for i in range(len(var.dimensions)):
                                var_time_step_dimensions[var.dimensions[i]]=var_time_step.shape[i]
                            if lDebugFlag: 
                                print var_time_step_dimensions    
                                print "\noutput variable attributes"
                                print var_atts
                                print "\noutput variable coordinate variables"
                                print var.coordinates
                                print "\noutput variable coordinate attribute variables:"
                            if lDebugFlag: 
                                for att in var_atts.keys():
                                    print att
                                    print var_atts["coordinates"]     

# open output file          fname_out used for testing, fname_out1 from config file

                            fname_out1 = output_dir + "nemo_" + (fname.split("/")[-1]).split(".")[0] + "_" + vname + "_timestep_" + str(time_step) + ".nc"
                            print "extended filename:", fname_out1
                            fname_out = "../nemo_split_test/test_split_" + str(time_step) + "_" + vname + ".nc"
                            print "writing to "+ fname_out1
                            try:
                               dsout = Dataset(fname_out1, "w", format = "NETCDF4_CLASSIC") #"NETCDF3_CLASSIC")
                            except:
                                print "Error creating output file ", fname_out
                                print "exiting"
                                sys.exit()

#create dimensions

                            if lDebugFlag:
                                print var.shape
                                print var_time_step_dimensions
                            for dname, the_dim in var_time_step_dimensions.items():
                               if dname == "time_counter":
                                   dname = "time"
                               if lDebugFlag:
                                   print "creating dimension", dname, the_dim
                               if dname == "time":
                                   dsout.createDimension(dname, size = 0)
                               else:
                                   dsout.createDimension(dname, the_dim) 

# create bounds

                            for dname, the_dim in file_dims.items():
                               if dname == "tbnds":
                                   dsout.createDimension(dname,the_dim)
                            
#get list of all required coordinate variables -from variable itself and coordinate attribute
     
                            all_coords_list=[]
                            bounds_list=[]
                            for v in var.coordinates.split():
                                all_coords_list.append(v)

                            for att in var_atts.keys():

                                if att == "coordinates":
                                    if lDebugFlag:
                                        print "variable coord atts:", 
                                    for item in var_atts["coordinates"].split():
                                        if item in all_coords_list:
                                            if lDebugFlag:
                                                print item, "already in coordinate list"
                                            continue
                                        else:
                                            if lDebugFlag:
                                                print item, "to be appended"
                                            all_coords_list.append(item)                                                             

                            if lDebugFlag:
                                print all_coords_list,len(all_coords_list)

# write out coordinate variable and attributes; amend name of time coordinate

                            for v_coord in all_coords_list:
                                var_coord = fin.variables[v_coord]
                                var_coord_time_dim =[]
                                for dim in var_coord.dimensions:
                                    if "time_counter" in dim:
                                        dim = "time"
                                    var_coord_time_dim.append(dim)
                                  
                                if v_coord == "time_counter":
                                    v_coord_time = "time"                                    
                                    out_var = dsout.createVariable(v_coord_time,var_coord.dtype, var_coord_time_dim)
                                else:
                                    out_var = dsout.createVariable(v_coord,var_coord.dtype, var_coord_time_dim)
                                if lDebugFlag:
                                    print v_coord,var_coord.dtype, var_coord.dimensions
                                    print type(out_var)

                                if "time" in v_coord:
                                    if lDebugFlag:
                                        print "tstep ***", time_step_value
                                    dtime = num2date(time_step_value,var_coord.units)
                                    out_var[:]=time_step_value
                                else:
                                    out_var[:] = var_coord[:]
                                if lDebugFlag:
                                    print "var_coord attributes: ",var_coord.ncattrs()

                                for ncattr in var_coord.ncattrs():
                                    if ncattr == "_FillValue":  
                                        f_val = getattr(var_coord, "_FillValue")
                                        out_var.setncattr(ncattr, fval) 

                                    if ncattr == "title": 
                                        continue

# Check for bounds variable **** need to get correct value ******

                                    if ncattr == "bounds":
                                        if lDebugFlag:
                                            print "bounds att for ", var_coord.getncattr(ncattr)
   
                                        bounds_list.append(var_coord.getncattr(ncattr)) 
                                                                                                   
#
# Setting attributes for required co-ordinate variable attribute: 
#
#                                    print dir(out_var),type(out_var)

                                    if v_coord == "time_counter" and var_coord.getncattr(ncattr) == "time_counter_bnds":
                                        print "var_coord.getncattr(ncattr)",var_coord.getncattr(ncattr)
                                        out_var.setncattr(ncattr, "time_bounds")
                                    elif v_coord == "time_counter" and ncattr == "units":
                                        time_units = var_coord.getncattr(ncattr)
                                        print time_units
                                        new_time_units=adjust_time(fname,time_units)
                                        out_var.setncattr(ncattr, new_time_units)
                                    else:    
                                        out_var.setncattr(ncattr, var_coord.getncattr(ncattr))                            

# adjust time to allow for non-uniform NEMO data start

# write bounds variable to output file

                            if lDebugFlag: 
                                print "bounds list", bounds_list
                            for bvar in bounds_list:
                                bound_var=fin.variables[bvar]
                            bound_var_time_dim=[]
                            for dim in bound_var.dimensions:
                                if dim == 'time_counter':
                                    dim = "time"
                                bound_var_time_dim.append(dim)
                            if lDebugFlag:
                               print "bounds"
                               print "bound_var ", bound_var
                               print "bound_var_dtype ",bound_var.dtype
                               print "bound_var_dimensions ", bound_var_time_dim
                               print "bound_var_ncattrs ", bound_var.ncattrs()
                            out_var = dsout.createVariable("time_bounds",bound_var.dtype, bound_var_time_dim)
                            out_var[:] = bound_var[:]
                            for ncattr in bound_var.ncattrs():
                                out_var.setncattr(ncattr, bound_var.getncattr(ncattr))                           
                            if lDebugFlag:
                                print dir(var_time_step),type(var_time_step)
                                print vname,var_time_step.dtype,var.dimensions
                            var_time_dim=[]
                            for dim in var.dimensions:
                                if dim == 'time_counter':
                                    dim = "time"
                                var_time_dim.append(dim)
                                                        
# write variable data        

                            if ncattr == "_FillValue":
                                f_val = getattr(varin, "_FillValue")
                                out_var = dsout.createVariable(vname,var_time_step.dtype, var_time_dim, fill_value = f_val)
                            else:
                                out_var = dsout.createVariable(vname,var_time_step.dtype, var_time_dim)
                            out_var[:] = var_time_step[:]
                                
# write variable attributes

                            var_att_list = modify_variable_attributes(var)

# Amend units string if required

                            if vname == "votemper":
                                if var_att_list["units"] == "C":
                                    var_att_list["units"] = "degC"
                            if vname == "sosaflx":
                                if var_att_list["units"] == "Kg/m2/s":
                                    var_att_list["units"] = "Kg/m^2/s"
                            if vname == "soicecov":
                                if var_att_list["units"] == "[0,1]":
                                    var_att_list["units"] = "1"

                            if vname == "soicealb":
                                if var_att_list["units"] == "[0,1]":
                                    var_att_list["units"] = "1"

                            if vname == "sowafldp":
                                if var_att_list["units"] == "Kg/m2/s":
                                    var_att_list["units"] = "Kg m-2 s-1"

                            if vname == "sowaflup":
                                if var_att_list["units"] == "Kg/m2/s":
                                    var_att_list["units"] = "Kg m-2 s-1"

                            if vname == "sosaflup":
                                if var_att_list["units"] == "Kg/m2/s":
                                    var_att_list["units"] = "Kg m-2 s-1"

# write MARS specific attributes amending elements fo thos calculated from file

                            cdftime=netcdftime.utime(dsout.variables["time"].units)
                            start_date_time=str(cdftime.num2date(float(dsout.variables["time_bounds"][0][0]))).split()
                            m_date=start_date_time[0].replace("-","")
                            m_time="".join(start_date_time[1].split(":")[:-1])
                            mars_cfg_dict["mars_date"] = m_date
                            mars_cfg_dict["mars_time"] = m_time

# need to adjust mars_level and levelist as required
                            #mars_cfg_dict["mars_level"] = ""
                            #mars_cfg_dict["mars_levelist"] =""                        

                            add_mars_atts(dsout, mars_cfg_dict, vname)
                            
                            for v in var_att_list:    
                                if lDebugFlag:
                                    print "* ",v, var_att_list[v]
                                if v == "_FillValue":
                                    continue
                                if v == "coordinates":
                                
                                    out_var.setncattr(v, var_att_list[v].replace("time_counter","time"))
                                else:
                                    out_var.setncattr(v, var_att_list[v])

                            cf_name=var_cfg_dict[vname][0]
                            print "CF Name is ",cf_name
                            out_var.setncattr("standard_name",cf_name)
                            out_var.setncattr("NEMO_variable_name",vname)

#write global attributes
                            if lDebugFlag:
                                print len(file_atts)
                            if len(file_atts) > 0:
                               for ncattr in file_atts:
                                   if ncattr == "history":
                                      continue
                                   dsout.setncattr(ncattr, file_atts[ncattr])

# write global attributes from config file

                            for key in at_cfg_dict:
                                dsout.setncattr(key, at_cfg_dict[key])   
                            
                            var_count = var_count + 1     
                                   
# close output file 
                            dsout.close()

#close input file 

        fin.close()
        
        print "Processed: " , fname 
        print "No. of Variables in file: " , len(fin_var_list)
        print "No. of Variables processed: ", var_count 
    
    
    
if __name__ == "__main__":
    main()

