#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# United States Department of Commerce
# NOAA (National Oceanic and Atmospheric Administration)
# National Weather Service
# Office of Water Prediction
#
# Authors:
#    Xuning Tan, Greg Fall	Created 02/16/2018
#
""" For mountain mapper down scale  
"""
__version__ = '1.0'

# Import System libraries
import time
import os
import logging
import argparse
import datetime 
import calendar
import numpy as np
import numpy.ma as ma
from netCDF4 import Dataset

# Default settings
# Name of the output log file
DEFAULT_LOG_FILENAME = "/tmp/mountain_mapper_downscale.log"

#  Format for logged messages
DEFAULT_LOG_FORMAT = "%(asctime)s : file - %(filename)s, " \
        "function: %(funcName)s, line: %(lineno)s, level: %(levelname)s - "\
        "%(message)s"

#  Format for the dates used in logged messages
DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S%z"

#  Level of messages that are written to the log
DEFAULT_LEVEL = "DEBUG"

#  Whether need to apply Mountain Mapper downscaling 
DOWNSCALE = False 

start = time.time()
def elapsed():
    return time.time() - start

#
def setup_arguments():
    """ Setup command line argument handling
    """
    parser = argparse.ArgumentParser(description='Mountain Mapper downscaling for National Water Model:')

    parser.add_argument( "--logfile", 
			  help="Specify log file name. The default log file is: " \
                          + DEFAULT_LOG_FILENAME,
			  type=str,
			  default=DEFAULT_LOG_FILENAME )
    parser.add_argument( "-v",
			 "--debug",
                         help="Display verbose messages.",
                         action="store_true" )
    parser.add_argument( "--version",
                         help="Display version number.",
                         action="store_true" )
    parser.add_argument( "--precipfile",
                         help="Specify the pathname for the RAINRATE netCDF file."
                              + " Or indicated in environment variable 'outFile'",
              		 type=str )
    parser.add_argument( "--PRISMdir",
                         help="Specify the pathname for the directory holds PRISM grids.",
              		 type=str )
    parser.add_argument( "--precipST",
                         help="Specify the start time of the current TIME STEP.",
              		 type=str )
    parser.add_argument( "--downscale",
                         help="Apply Mountain-Mapper downscaling for NWM",
              		 action="store_true" )
    return parser.parse_args()

def setup_logging( commandline_args ):
    """ Setup run-time logging.

    Perform logging method initialization for the application.

    """

    log_level = None
    if commandline_args.debug:
       log_level = logging.DEBUG
    else:
       log_level = logging.INFO

    logging.basicConfig( filename=commandline_args.logfile,
			 filemode='w',
                         level=log_level,
                         format=DEFAULT_LOG_FORMAT,
                         datefmt=DEFAULT_DATE_FORMAT )

    logging.debug( "The logging method has been setup." )

def main():

    x = y = numx = numy = denx = deny = 0

    """ Run the program.
    """
    # Get command line arguments.
    args = setup_arguments()

    # Set logging method.
    setup_logging( args )

    if args.version:
        print( "VERSION " + str( __version__ ) )
        return
    
    # Check for alternative precip file specification, 
    # Command line setting overrides environment variables
    # Refer to the env variables setting in nwm.v1.2.0/ush/down/rh_correction_narr.ncl
    if args.precipfile:
        precip_file = args.precipfile
        logging.debug( "Command line 'precipfile' of Mountain Mapper: " + precip_file )
    elif ( os.environ.has_key ('outFile') ): 
        precip_file = os.environ.get( 'outFile' ) 
        logging.debug( "Env variable 'outFile': " + precip_file )
    else:
        logging.error( "No pathname specified for the 'precipfile' of Mountain Mapper" )
        print( "ERROR: No pathname specified for the 'precipfile' of Mountain Mapper" )
        return

    # Check the directory specified for PRISM grids files
    if args.PRISMdir:
        PRISM_dir = args.PRISMdir
        logging.debug( "Command line 'PRISMdir' of Mountain Mapper: " + PRISM_dir )
    else:
        logging.error( "No pathname specified for the 'PRISMdir' of Mountain Mapper" )
        print( "ERROR: No pathname specified for the 'PRISMdir' of Mountain Mapper" )
        return

    # Check whether precipST has been passed on 
    if args.precipST:
        precip_accum_start_time = args.precipST
        logging.debug( "Command line 'precipST' was set to: " + precip_accum_start_time )
    else:
        logging.error( "No precipST specified" )
        print( "ERROR: No precipsST specified" )
        return

    # Check whether Mountain-Mapper downscaling has been requested on command line
    if args.downscale: 
        DOWNSCALE = True 
        logging.debug( "Command line 'downscale' was set to True" )
        # Otherwise check the envirnment variable'DOWNSCALE_RAINRATE'
    elif os.environ.has_key ('DOWNSCALE_RAINRATE'):
        DOWNSCALE = os.environ.get( 'DOWNSCALE_RAINRATE')
        logging.debug( "'DOWNSCALE_RAINRATE' environment variable is: " + DOWNSCALE )
    else:
        logging.error( "Neither '--downscale' argument nor 'DOWNSCALE_RAINRATE' environment variable has been set" ) 
        print( "ERROR: Neither '--downscale' argument nor 'DOWNSCALE_RAINRATE' environment variable has been set" ) 
        return

    # Check whether need to run Mountain-Mapper, if not, exit this program
    if DOWNSCALE:
        logging.info( "Prepare for Mountain Mapper Downscale..." )
        print( "Prepare for Mountain Mapper Downscale..." )
    else:
        logging.info( "Not Apply Mountain Mapper Downscale, bye..." )
        print( "Not Apply Mountain Mapper Downscale, bye..." )
        return

    # Open precip file
    try:
        precip_fid = Dataset(precip_file, 'r+')
    except IOError:
        logging.error( "Cannot open " + precip_file )
        print( "ERROR: Cannot open " + precip_file )
        return
    
    # Get dimensions from input netcdf file
    for name, dimension in precip_fid.dimensions.items():
        if name == 'x':
            x = len(dimension)
        if name == 'y':
            y = len(dimension)
    # y = 3840; x = 4608
    logging.debug( "numx = " + str(numx) + "; numy = " + str(numy))
    # Get RAINRATE data from input netcdf file
    for var in precip_fid.variables:
        if var == 'RAINRATE':
            rainrate_data = precip_fid.variables[var][0,:,:]
            break
    if var != 'RAINRATE':
        logging.error( "No 'RAINRATE' variable found in file " + precip_file )
        print( "ERROR: No 'RAINRATE' variable found in file " + precip_file )
        precip_fid.close()
        return

    # Retrieve the calendar month from the start of the current TIME STEP"
    # Assume the time format should be 'YYYYMMDDHH' 
    logging.debug( "'precip_accum_start_time' is " + precip_accum_start_time )
    cal_month = calendar.month_abbr[int(precip_accum_start_time[4:6])]
    logging.debug( "'Abbreviate Calendar month' is " + cal_month )
    
    # Test here
    #cal_month = "Jan"
 
    # Check the PRISM grids
    numerator_file = PRISM_dir + "/PRISM_Precip_Clim_" + cal_month + "_NWM_Mtn_Mapper_Numer.nc"
    denominator_file = PRISM_dir + "/PRISM_Precip_Clim_" + cal_month + "_NWM_Mtn_Mapper_Denom.nc"

    # Open numerator file
    try:
        numerator_fid = Dataset(numerator_file, 'r')
    except IOError:
        logging.error( "Cannot open " + numerator_file )
        print( "ERROR: Cannot open " + numerator_file )
        return

    # Get dimensions from numerator file
    for name, dimension in numerator_fid.dimensions.items():
        if name == 'x':
            numx = len(dimension)
        if name == 'y':
            numy = len(dimension)
    # y = 3840; x = 4608
    logging.debug( "numx = " + str(numx) + "; numy = " + str(numy))
    # Get climatological precipitation data from numerator file
    for var in numerator_fid.variables:
        if var == 'Data':
            numerator_data = numerator_fid.variables[var][:]
            break
    if var != 'Data':
        logging.error( "No 'Data' variable found in file " + numerator_file )
        print( "ERROR: No 'Data' variable found in file " + numerator_file )
        numerator_fid.close()
        return

    numerator_fid.close()

    # Open denominator file
    try:
        denominator_fid = Dataset(denominator_file, 'r')
    except IOError:
        logging.error( "Cannot open " + denominator_file )
        print( "ERROR: Cannot open " + denominator_file )
        return

    # Get dimensions from denominator file
    for name, dimension in denominator_fid.dimensions.items():
        if name == 'x':
            denx = len(dimension)
        if name == 'y':
            deny = len(dimension)
    # y = 3840; x = 4608
    logging.debug( "denx = " + str(denx) + "; deny = " + str(deny))
    # Get climatological precipitation data from denominator file
    for var in denominator_fid.variables:
        if var == 'Data':
            denominator_data = denominator_fid.variables[var][:]
            break
    if var != 'Data':
        logging.error( "No 'Data' variable found in file " + denominator_file )
        print( "ERROR: No 'Data' variable found in file " + denominator_file )
        denominator_fid.close()
        return
        
    denominator_fid.close()

    if not ( x == numx == denx != 0 ):
        logging.error( "Dimension x not match in precip file, numerator or denominater files" )
        return
    if not ( y == numy == deny != 0 ):
        logging.error( "Dimension y not match in precip file, numerator or denominater files" )
        return

    # Apply NWP Mountain Mapper
    logging.info( "Applying NWP Mountain Mapper computation, please wait...")
    print( "Applying NWP Mountain Mapper computation, please wait...")
    ma.masked_where(denominator_data <= 0.01, denominator_data)
    rainrate_data *= numerator_data
    rainrate_data /= denominator_data
    logging.info( "Done with NWP Mountain Mapper computation.")
    print( "Done with NWP Mountain Mapper computation.")
    print ('It took %0.3f s' % elapsed())

    # Replace the RAINRATE data with the Mountain Mapper Downscaled data
    precip_fid.variables['RAINRATE'][0,:,:] = rainrate_data
    precip_fid.close()
    print( "Done updating the precipfile/outfile.")
    print ('It took %0.3f s total' % elapsed())

# Module run directly as a script 
if __name__ == '__main__':
    main()