#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -*- 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
import calendar

def seconds_in_month(month, year):
    nomatter, daysinmonth = calendar.monthrange(year, month)
    return daysinmonth * 24 * 60 * 60


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


print "to run: "
print "split monthly nemo data file to extract 1 variable (salinity) " +\
"-i <input file> -o <output file prefix> "
print "to Run "
print "python split_monthly.py "
print "-i <nemo monthly file>" 
print "-o <nemo monthly split file> "


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 == '-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



fname_out_temp = output_filename
fname = input_filename

#flist = glob.glob("/hugetmp/cera20c/model_data/2366_1m_1900*.nc")
#print flist

islDebugFlag = 0

#cmd = "rm -f /hugetmp/cera20c/model_data/split_month/*.nc*"
#os.system(cmd)

#op_dir = "/hugetmp/cera20c/model_data/split_month/"

co_var = "vosaline"

for var_count in [1]:
       fin = Dataset(fname, 'r')
     
       #fname_out_temp = op_dir + co_var + "_" + fname.split("/")[-1] + "4"
       
       print "writing INTERIM NC4 file to " + fname_out_temp

       try:
          dsout = Dataset(fname_out_temp, "w", \
          format="NETCDF4_CLASSIC")
       except:
          print "Error creating output file ", fname_out_temp
          print "exiting"
          sys.exit()
       date_str = fname.split("/")[-1].split("_")[2]
#       print "**",date_str

       year=date_str[0:4]
       month= date_str[4:6]
       day=date_str[6:9].lstrip("0")
       print year, month,day
       time_units_str = "seconds since " + year +"-"+ month+"-"+ date_str[6:9] +" 00:00:00"
       print time_units_str       
       

       var_out_dimensions = {}
       var_out = fin.variables[co_var] 

#       print var_out.shape
#       print var_out
#       print "dimesniosns", var_out.dimensions


# set dimensions

       dim_count = 0

       for dim_name in var_out.dimensions:
           if dim_name == "time_counter":
               var_out_dimensions["time"] = var_out.shape[dim_count]
           elif dim_name == "deptht":
               var_out_dimensions["depth"] = var_out.shape[dim_count]
           else:    
               var_out_dimensions[dim_name] = var_out.shape[dim_count]
           dim_count = dim_count + 1
           print dim_name, var_out_dimensions
       if islDebugFlag:
           print var_out_dimensions
           print var_out.shape
           print type(var_out), type(var_out)
       
       print var_out_dimensions
       print fin.variables[co_var].datatype
       print 'dimensions',var_out_dimensions
       for dname, the_dim in var_out_dimensions.iteritems():
           if dname == "time_counter":
               dsout.createDimension(dname, size=0)
           else:
               dsout.createDimension(dname, the_dim)    

# get time values from bounds

       intyear=int(year)
       intmonth=int(month)
       secs=seconds_in_month(intmonth, intyear)
       
       time_set = 0
       lower_time_bound = 0
       upper_time_bound = secs
       
#       dstime_bounds'][0,0] = 0
#       dset.variables['time_bounds'][0,1] = secs 
               
# create variable
   
       outVar = dsout.createVariable(co_var, \
       fin.variables[co_var].datatype, \
       ("time", "depth", "y", "x"))

      
       for k in fin.variables[co_var].ncattrs():
           print k
           if k == "coordinates": 
              print "reset coords *****",k, type(fin.variables[co_var].getncattr(k)), 
              print type("time depth nav_lat nav_lon") 
              co_str = "time depth nav_lat nav_lon"
              outVar.setncattr("coordinates",co_str)
           else: 
               outVar.setncattr(k,fin.variables[co_var].getncattr(k))
           
       outVar[:] = fin.variables[co_var][:] 
       
# add extra attributes
       ec_var_atts = {"standard_name" : "sea_water_salinity", \
                "cell_methods":"time: mean (interval: 1.0 month)" , \
                "mars_stream" : "enda" , \
                "mars_class" : "ep" , \
                "mars_type" : "an" , \
                "mars_expver" : 2366 , \
                "mars_levtype" : "dp" , \
                "mars_time" : 0 , \
                "mars_date" : int(date_str) , \
                "mars_param" : 34 , \
                "mars_levelist" : 1 , \
                "mars_number" : 1 }
       
       for ncattr in ec_var_atts.keys():
           outVar.setncattr(ncattr, ec_var_atts[ncattr])        
       
          
# create time variable 
       
       outVar = dsout.createVariable("time", \
       fin.variables[co_var].datatype, \
       ("time",))
       
       for k in fin.variables["time_counter"].ncattrs():
           if k == "time_origin": 
               continue
           outVar.setncattr(k,fin.variables["time_counter"].getncattr(k))
           if k == "bounds": 
               outVar.setncattr(k,"time_bnds")  
           if k == "units":
               outVar.setncattr(k,time_units_str)  
                    

       outVar[:] = time_set

#       outVar[:] = fin.variables["time_counter"][:]        #for original values

# create depth variable
       
       outVar = dsout.createVariable("depth", \
       fin.variables["deptht"].datatype, \
       ("depth",))
       for k in fin.variables["deptht"].ncattrs():
           if k == "title": 
               outVar.setncattr("title","depth")
           else:
               outVar.setncattr(k,fin.variables["deptht"].getncattr(k))

       outVar[:] = fin.variables["deptht"][:]    

#nav lon
                 
       outVar = dsout.createVariable("nav_lon", \
       fin.variables["nav_lon"].datatype, \
       ("y","x"))
       outVar.setncatts({k: \
       fin.variables["nav_lon"].getncattr(k) \
       for k in fin.variables["nav_lon"].ncattrs()})

       outVar[:] = fin.variables["nav_lon"][:]                    

#nav lat
       
       outVar = dsout.createVariable("nav_lat", \
       fin.variables["nav_lat"].datatype, \
       ("y","x"))
       outVar.setncatts({k: \
       fin.variables["nav_lat"].getncattr(k) \
       for k in fin.variables["nav_lat"].ncattrs()})

       outVar[:] = fin.variables["nav_lat"][:]     

# time bounds

       tbnds_length = 2 
       dsout.createDimension("tbnds",tbnds_length )
       outVar = dsout.createVariable("time_bnds", \
       fin.variables["time_counter_bnds"].datatype, \
       ("time", "tbnds"))
       outVar.setncatts({k: \
       fin.variables["time_counter_bnds"].getncattr(k) \
       for k in fin.variables["time_counter_bnds"].ncattrs()})
       outVar[0,0] = lower_time_bound
       outVar[0,1] = upper_time_bound
       
#       outVar[:] = fin.variables["time_counter_bnds"][:] #for original values
       
       
       
# Global attributes       
       
#      
       nc_attrs_dict = {}
       nc_atts_list = fin.ncattrs()
       print nc_atts_list, type(nc_atts_list)
       print "NetCDF Global Attributes:"

       if len(nc_atts_list) == 0:
           print "No Global attributes in source file"

       for nc_attr in nc_atts_list:
           print nc_attr
           nc_attrs_dict[nc_attr] = fin.getncattr(nc_attr)       
       print nc_attrs_dict
           
           
       for ncattr in nc_attrs_dict.keys():
           if ncattr == "history" or ncattr == "nco_openmp_thread_number":
               continue
           print "***",ncattr, nc_attrs_dict[ncattr]
           dsout.setncattr(ncattr, nc_attrs_dict[ncattr])       

       ec_global_atts={"comment":"Produced at ECMWF",
       "title":"NEMO model output",
       "Conventions" : "CF-1.6",
       "source":"NEMO V3.4" ,
       "references": "Madec G. 2008: NEMO ocean engine Note du Pole de " + \
       "modélisation, Institut Pierre-Simon Laplace (IPSL), France, No 27 "+ \
       "ISSN No 1288-1619. \\nhttp://www.nemo-ocean.eu/content/download/21612/"+ \
       "97924/file/NEMO_book_3_4.pdf",
       "institution":"ECMWF"}
        
       for ncattr in ec_global_atts.keys():

           print "***",ncattr, ec_global_atts[ncattr]
           dsout.setncattr(ncattr, ec_global_atts[ncattr])       
        
       dsout.close()
     
       fin.close()