#!/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
#
# 
# v3 modified from V2 to take command line input file, output file (prefix),
# type of file to be processed and configuration file to amend attributes etc.
#
# Currently only works for NEMO 3hr and 1D data files - Not 1 month/ non-NEMO data
#
# K Marsh 8/8/16
#
#*****************************************************************
# to run:
# python convert_nemo_to_1_variable_all_levels_one_timestep_v2.py 
# -i copy_of_2366_1d_1900022721_1900022821_grid_T_02.nc 
# -o copy_of_2366_1d_1900022721_1900022821_grid_T_02 
# -k NEMO -c convert_nemo_to_1_variable_all_levels_one_timestep.cfg
# 
# Note that this appends the variable name, timestep and "INTERIM.nc3" 
# to the output filename
#
from netCDF4 import Dataset, netcdftime, num2date
import glob
import sys
import getopt
import datetime
import os

def process_config_file(cfg_file_in):
    islDebugFlag = 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("=")
            op = option.strip()
            v = value.strip()
            os.environ[op] = "$" + op + ":" + v
        else:
            if islDebugFlag:
                print line
            option, value = line.split("=")
            op = option.strip()
            v = value.strip()
            if op == "db_name":
                db_name = v
            elif op == "host_name":
                host_name = v
            elif op == "user_name":
                user_name = v
            elif op == "pass_name":
                pass_name = v
            elif op == "mars_type":
                mars_dict["mars_type"] = v
            elif op == "mars_class":
                mars_dict["mars_class"] = v
            elif op == "mars_stream":
                mars_dict["mars_stream"] = v
            elif op == "mars_expver":
                mars_dict["mars_expver"] = v
            elif op == "mars_date":
                mars_dict["mars_date"] = v
            elif op == "mars_time":
                mars_dict["mars_time"] = v
            elif op == "mars_step":
                mars_dict["mars_step"] = v
            elif op == "mars_number":
                mars_dict["mars_number"] = v
            elif op == "mars_levelist":
                mars_dict["mars_levelist"] = v
            elif op == "mars_levtype":
                mars_dict["mars_levtype"] = v
            elif op == "mars_param":
                mars_dict["mars_param"] = v
            elif op == "fnames":
                fnames.append(v)
            elif op == "input_dir":
                fnames = glob.glob(v)
            elif op == "output_dir":
                op_dir = v
            elif op == "global_attributes_Conventions":
                at_dict["Conventions"] = v
            elif op == "global_attributes_title":
                at_dict["title"] = v
            elif op == "global_attributes_institution":
                at_dict["institution"] = v
            elif op == "global_attributes_references":
                at_dict["references"] = v
            elif op == "global_attributes_source":
                at_dict["source"] = v
            elif op == "global_attributes_history":
                at_dict["history"] = v
            elif op == "global_attributes_comment":
                at_dict["comment"] = v
            elif op == "nemo_var_name":
                var_dict[v.split(":")[0]] = \
                [v.split(":")[1], v.split(":")[2], v.split(":")[3], \
                v.split(":")[4], v.split(":")[5]]
            else:
                print "Warning: Unmatched option", op
    cfg_file_in.close()



    if not op_dir.endswith("\/"):
        op_dir = op_dir + "/"
    if len(mars_dict) != 11:
        print "Missing Mars elements: exiting"
        sys.exit()
    if islDebugFlag:
        print "MARS dictionary is :", mars_dict
    return at_dict, mars_dict, var_dict

def get_original_dimensions_and_global_attributes(file_in):
    islDebugFlag0 = True
#
#get dimensions
#
    dim_list = file_in.dimensions
    dim_dict = {}
    nc_attrs = {}
#
#Copy dimensions from whole file
#
    for dname, the_dim in file_in.dimensions.iteritems():
        dim_dict[dname] = len(the_dim)
#
# get global attributes
#
    nc_atts = file_in.ncattrs()
    print nc_atts, type(nc_atts)
    print "NetCDF Global Attributes:"
    print dir(file_in)
    if len(nc_atts) == 0:
        print "No Global attributes in source file"
    if islDebugFlag0:
        for nc_attr in nc_atts:
            print nc_attr
            nc_attrs[nc_attr] = file_in.getncattr(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):

    islDebugFlag1 = 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 islDebugFlag1:
            print ncattr
            print var_in.getncattr(ncattr)
        var_atts_out[ncattr] = var_in.getncattr(ncattr)
    if islDebugFlag1:
        print "variable attributes", var_atts_out
    return var_atts_out

def add_mars_atts(fid, att_dict, var):
    islDebugFlag2 = False
    for item_att in att_dict.keys():
        if islDebugFlag2:   
            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):
    islDebugFlag3 = False
    filename_time = (fname_in.split("/")[-1]).split("_")
    time1 = filename_time[2][-2:]
    time2 = filename_time[3][-2:]

    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 islDebugFlag3:
        print time_units_new

#        sys.exit(1)
    return time_units_new
    
#
# NEMO 1 month averages do not have the correct date/units set.
# These default to the filename and should be set to the first day of each month
# votemper_2366_1m_1900_grid_T_02.nc
# 2366_1d_1900022821_1900030121_grid_T_02.nc
# 2366_3h_1900022821_1900030121_grid_T_03.nc



def add_nemo_1month_date_values(fname_in):
    day_number_list = []
    dates_list = []
    cpts =  fname_in.split("_")
    if cpts[2] == "1m":
        yr_str=cpts[3]
        print yr_str
    else:
        print "wrong 1 month filename; exit"
        sys.exit()
    units_str = "days since "+ yr_str + "-1-1 00:00:00 " 
    print units_str
    dy_str = "01"
    for mn in range(1,13):
        mn_str = str(mn).zfill(2)
        print yr_str, mn_str,dy_str
        date_str = yr_str + "." + mn_str +"."+ dy_str
        fmt = '%Y.%m.%d'
        dt = datetime.datetime.strptime(date_str, fmt)
        print dt
        tt = dt.timetuple()
        day_number_list.append(tt.tm_yday -1)
        dates_list.append(date_str.replace(".",""))
    print day_number_list
    print dates_list


def main():
    add_nemo_1month_date_values("votemper_2366_1m_1900_grid_T_02.nc")

    islDebugFlag = True
    hasConfigFile = False
    hasInputFile = False
    hasOutputFilePath = False
    hasDataType = False 
    print "to run: "
    print "convert_nemo_to_1_variable_all_levels_one_timestep.py " +\
    "-i <input file> -o <output file prefix> -k <data type> -c <config file>"
    print "to Run "
    print "python convert_nemo_to_1_variable_all_levels_one_timestep_v2.py "
    print "-i copy_of_2366_1d_1900022721_1900022821_grid_T_02.nc" 
    print "-o copy_of_2366_1d_1900022721_1900022821_grid_T_02" 
    print "-k NEMO -c convert_nemo_to_1_variable_all_levels_one_timestep.cfg"
 
    print "Note that this appends the variable name, timestep and <INTERIM.nc3>" 
    print "to the output filename"

    print "Parsing input config file"
    options, remainder = getopt.getopt(sys.argv[1:], 'c:i:o:k:')
    print options, remainder
    if len(options) == 0:
        print "Option required; Exiting"
        sys.exit()
    for opt, arg in options:
        print opt, arg
        if opt == '-c':
            cf_cfg_filename = arg
            if not os.path.isfile(cf_cfg_filename):
                print "config file does not exist; exiting"
                sys.exit()
            print "Config file is ", cf_cfg_filename
            hasConfigFile = True
        if opt == '-i':
            input_filename = arg
            if not os.path.isfile(input_filename):
                print "Input file does not exist; exiting"
                sys.exit()
            print "Input file is ", input_filename   
            hasInputFile = True         
        if opt == '-o':   
            output_filename = arg
            if os.path.isfile(output_filename):
                print "Output file does exist and will be overwritten"
            print "Output file is ", output_filename            
            hasOutputFilePath = True            
        if opt == '-k':
            file_type = arg
            if not file_type == 'NEMO':
                print "Need file type to be set (e.g. NEMO)"
                sys.exit()
            print "file type is NEMO"
            hasDataType = True 
        #else:
        #    print "incorrect command line arguement; Exiting"
        #    sys.exit()
    try:
        cfg_file = open(cf_cfg_filename, "r")
    except:
        print "Configuration file to process does not exit; Exiting"


#  Initialise variables to hold file information


    files = []
    at_cfg_dict = {}
    mars_cfg_dict = {}
    var_cfg_dict = {}
    start_date = ""
    start_time = ""
    
    files.append(input_filename)
    print len(files)
    print files
    at_cfg_dict, mars_cfg_dict, var_cfg_dict = \
    process_config_file(cfg_file)

    if islDebugFlag:
        print "at_cfg_dict", at_cfg_dict
        print mars_cfg_dict
        print var_cfg_dict
        print "files to process: "
        print files
        print "\n"

#files = ["/hugetmp/cera20c/split/2366_1d_1900022721_1900022821_grid_T_02.nc",]
    print files, "**"
    if len(files) == 0:
        print "No files to process - exiting !"
        sys.exit()

    for fname in files:
        is3HourFile = False
        is1DayFile = False
        is1MonthFile = False
        fin = Dataset(fname, 'r')
        if "_3hr_" in fname:
            is3HourFile = True
        if "_1m_" in fname:
            is1MonthFile = True
        if "_1d_" in fname:
            is1DayFile = True
        
        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 islDebugFlag:
                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 == "so14chgt": #temp var to select -use below in operations

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

                for time_step in range(var.shape[0]):
                    var_out_dimensions = {}
                    var_out = var[slice(time_step, time_step+1, None)]
                    dim_count = 0
                    for dim_name in var.dimensions:
                        var_out_dimensions[dim_name] = var_out.shape[dim_count]
                        dim_count = dim_count + 1
                    if islDebugFlag:
                        print time_step
                        print var_out_dimensions
                        print var_out.shape
                        print type(var), type(var_out)

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

                    fname_out = output_filename \
                    + "_" + vname + "_timestep_" + str(time_step) + ".nc"
                    fname_out_temp = fname_out.replace(".nc", "_INTERIM.nc3")
 
                    print "writing INTERIM NC3 file to " + fname_out_temp

                    try:
                        dsout = Dataset(fname_out_temp, "w", \
                        format="NETCDF3_CLASSIC")
                    except:
                        print "Error creating output file ", fname_out_temp
                        print "exiting"
                        sys.exit()

                    for dname, the_dim in var_out_dimensions.iteritems():
                        if dname == "time_counter":
                            dsout.createDimension(dname, size=0)
                        else:
                            dsout.createDimension(dname, the_dim)
                    if "tbnds" in file_dims:
                        dsout.createDimension("tbnds", file_dims["tbnds"])

                     
                        
                    coords = []
                    for co in var.coordinates.split():
                        coords.append(co)
                    for va in var.ncattrs():
                        if va == "coordinates":
                            va_co = var.getncattr(va).split()
                            for v in va_co:
                                if v in coords:
                                    continue
                                else:
                                    coords.append(va_co)

# if a monthly file, manually add time bounds
                    print "is1month file is ", is1MonthFile

                    if is1MonthFile:
                        dsout.createDimension("tbnds", 2)       
                        if "deptht" not in coords:
                            coords.insert(0,"deptht")
                        
                        if "time_counter" not in coords:
                            coords.insert(0,"time_counter")                        

#create time coordinate variables, attributes, set values

                    if islDebugFlag:
                        print "coords", coords
                    hasTimeBounds = False
                    for co_var in coords:
                        if islDebugFlag:
                            print "process", co_var
                        if co_var == "time_counter":
                            if "bounds" in fin.variables[co_var].ncattrs():
                                hasTimeBounds =  True
                                print "set time bounds"
                                time_bounds_var_name = \
                                fin.variables[co_var].getncattr("bounds")

                            outVar = dsout.createVariable("time", \
                            fin.variables[co_var].datatype, \
                            fin.variables[co_var].dimensions)

                            for k in fin.variables[co_var].ncattrs():
                                if k == "bounds":
                                    outVar.setncatts({k: "time_bnds"})
                                elif k == "units":
                                    time_units = \
                                    fin.variables[co_var].getncattr(k)
                                    if islDebugFlag:
                                        print "time units are ", time_units
                                    new_time_units = \
                                    adjust_time(fname, time_units)
                                    outVar.setncattr(k, new_time_units)
                                elif k == "time_origin":
                                    time_origin = \
                                    fin.variables[co_var].getncattr(k)
                                    new_time = \
                                    adjust_time(fname, time_origin)
                                    if islDebugFlag:
                                        print time_origin, new_time, time_origin
                                    new_time_origin_str = \
                                    time_origin.split()[0] + " " + new_time.split()[1]
                                    outVar.setncattr(k, new_time_origin_str)
                                else:
                                    outVar.setncatts({k: \
                                    fin.variables[co_var].getncattr(k)})
                            outVar[:] = fin.variables[co_var][time_step]

# create other coord variables
                        else:
                            outVar = dsout.createVariable(co_var, \
                            fin.variables[co_var].datatype, \
                            fin.variables[co_var].dimensions)
                            outVar.setncatts({k: \
                            fin.variables[co_var].getncattr(k) \
                            for k in fin.variables[co_var].ncattrs()})
                            outVar[:] = fin.variables[co_var][:]
# create time bounds
                    print "hasTimebounds  is ", hasTimeBounds
                    if hasTimeBounds:
                        outVar = dsout.createVariable("time_counter_bnds", \
                        fin.variables["time_counter_bnds"].datatype, \
                        fin.variables["time_counter_bnds"].dimensions)
                        outVar.setncatts({k: \
                        fin.variables["time_counter_bnds"].getncattr(k) \
                        for k in fin.variables["time_counter_bnds"].ncattrs()})

                        outVar[0, :] = \
                        fin.variables["time_counter_bnds"][time_step][0]
                        outVar[:, 1] = \
                        fin.variables["time_counter_bnds"][time_step][1]
                        
                    if is1MonthFile:
                        outVar = dsout.createVariable("time_counter_bnds", \
                        fin.variables["time_counter"].datatype, ("time_counter","tbnds"))
                        

# amend variable attributes

                    var_att_list = modify_variable_attributes(var)
                    if islDebugFlag:
                        print var_att_list
                        print var_cfg_dict
                    layer_needed = False  
                                      
# add any additional coordinate information from the config file
# add scalar dimension and scalar coordinate variable (depth) if required for layer integral

                    if var_cfg_dict[vname][3] != "None":
                        scalar_cpts = var_cfg_dict[vname][3].split("_")
                        scalar_name = scalar_cpts[0]
                        scalar_units = scalar_cpts[2]
                        scalar_value = scalar_cpts[1]
                        scalar_long_name = var_cfg_dict[vname][4]
                        if islDebugFlag:
                            print scalar_cpts

                        if scalar_name == "depth":
                            if islDebugFlag:
                                print "found depth layer"
                            layer_needed = True
                            dsout.createDimension(scalar_name, size=1)
                            outvar=dsout.createVariable(scalar_name,'f4', )
                            outvar[:] = float(scalar_value)
                            outvar.units = scalar_units
                            outvar.positive = "down"
                            outvar.bounds = "depth_bnds"
                            outvar.long_name = scalar_long_name
                            var_att_list["coordinates"] = \
                            var_att_list["coordinates"].replace("time_counter", "time_counter "\
                             + scalar_name)
                            dsout.createDimension("dbnds", size=2)
                            outvar=dsout.createVariable("depth_bnds", 'f4', ("depth", "dbnds"))
                            outvar[0, 0] = 0 
                            outvar[0, 1] = float(scalar_value)

# add scalar variable if required for coordinate                            

                        if layer_needed is False:  
                            if "." in scalar_value:
                                dsout.createVariable(scalar_name, 'f4', )
                                scalar_var = dsout.variables[scalar_name]
                                scalar_var[:] = float(scalar_value)
                            else:
                                dsout.createVariable(scalar_name, 'i4', )
                                scalar_var = dsout.variables[scalar_name]
                                scalar_var[:] = int(scalar_value)
                            scalar_var.units = scalar_cpts[2]
                            scalar_var.long_name = scalar_long_name
                            var_att_list["coordinates"] = \
                            scalar_name + " " + var_att_list["coordinates"]

# Create variable, attributes, set values
                    if layer_needed:
                        if len(var.dimensions) == 3:
                            depth_layer_dims = (var.dimensions[0], "depth", \
                            var.dimensions[1], var.dimensions[2])
                        else:
                            print "dimension length problem for layer"
                            sys.exit()
                        outVar = \
                    dsout.createVariable(vname, var.datatype, depth_layer_dims)
                    else:
                        outVar = \
                        dsout.createVariable(vname, var.datatype, var.dimensions)

                    if "deptht" in var_att_list["coordinates"]:
                        print "found depth info for " + vname 
                        var_att_list["coordinates"] = \
                        var_att_list["coordinates"].replace("deptht", "depth")
                        
                    if "time_counter" in var_att_list["coordinates"]:
                        var_att_list["coordinates"] = \
                        var_att_list["coordinates"].replace("time_counter", \
                        "time")

                    for v in var_att_list:
                        if islDebugFlag:
                            print "* ", v, var_att_list[v]

                        outVar.setncattr(v, var_att_list[v])

# cf standard name
                    if len(var_cfg_dict[vname][0]) > 0:
                        cf_name = var_cfg_dict[vname][0]
                        outVar.setncattr("standard_name", cf_name)
                        outVar.setncattr("NEMO_variable_name", vname)
# cf standard units
                    if len(var_cfg_dict[vname][1]) > 0:
                        var_units = var_att_list["units"]
                        cf_units = var_cfg_dict[vname][1]
                        outVar.setncattr("units", cf_units)
                        outVar.setncattr("NEMO_units", var_units)

#  Add cell methods for variable
                    isCellMethodsBounds = False
                    if "1d" in fname:
                        outVar.setncattr("cell_methods", "time: mean (interval: 1.0 day)")
                    if "3h" in fname:
                        outVar.setncattr("cell_methods", "time: mean (interval: 3.0 hours)")                        
                    if isCellMethodsBounds:
                        outVar.setncattr("cell_methods", "time: mean (over time bounds)")

                    outVar[:] = var_out[:]
# write global attributes
                    if islDebugFlag:
                        print len(file_atts)
                    if len(file_atts) > 0:
                        print file_atts
                        for ncattr in file_atts:
                            if ncattr == "history":
                                continue
                            print "***",ncattr, file_atts[ncattr]
                            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])

# write MARS specific attributes amending elements to those calculated from file
                    print dsout.variables
                    cdftime = netcdftime.utime(dsout.variables["time"].units)
                    if hasTimeBounds:
                        start_date_time = \
                        str(cdftime.num2date\
                        (float(dsout.variables["time_counter_bnds"][0][0]))).split()
                    else:
                        start_date_time = str(cdftime.num2date\
                        (fin.variables["time_counter"][time_step])).split()
                        print start_date_time
                    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)
                    if "time_counter" in dsout.dimensions:
                        dsout.renameDimension('time_counter', 'time')
                    if "time_counter_bnds" in dsout.variables:
                        dsout.renameVariable('time_counter_bnds', 'time_bnds')
                    if "deptht" in dsout.dimensions:
                        dsout.renameDimension('deptht', 'depth')
                    if "deptht" in dsout.variables:
                        dsout.renameVariable('deptht', 'depth')
                    var_count = var_count + 1
                    dsout.close()
        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()
