# -*- coding: utf-8 -*-
#--------------------------------------------------------------------#
# Snow Amount Index                                                  #
# Filename: snow_amount.py                                           #
# Description: Generates impacts for snow amount and snow rates      #
# Authors: Nathan Foster/Mike Muccilli/Josh Kastman                  # 
# Change Log:                                                        #
# Output for rolling 24 hour periods - Kastman 2018-09-01            #
# New Snow Climatology - Kastman 2018-10-09                          #
# Updates to Snow Climatology - Kastman 2019-09-01                   #
# Updated for rolling 24 hr output - Kastman 2021-6-28               #
#--------------------------------------------------------------------#

# Imports

import arcpy, os, shutil, datetime, time, sys
from datetime import datetime
from time import gmtime, strftime
import logging
from datetime import datetime
import traceback
from shutil import copyfile
from arcpy import env
from arcpy.sa import *
from shutil import copyfile
import glob
import os.path
from os import path

import wssi_config as cfg
import wssi_general_functions as gfunc

# set up log file
log_stamp = strftime("%Y%m%d%H", gmtime())
logging.basicConfig(filename=cfg.log_dir+'/wssi_log_'+str(log_stamp)+'_snow.log',level=logging.DEBUG, format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

#currenthour = datetime.utcnow().strftime("%H")

def snow_amount():
    sn_amt_start = strftime("%a, %d %b %Y %H:%M:%S", gmtime())
    ##print ("WSSI Snow Amount Start Time "+sn_amt_start)

    #-----------------#
    # Set local lists #
    #-----------------#
    file = cfg.work_dir+'precp_num_list.txt'
    data = open(file, 'r')
    precip_num_text = data.readline()
    precip_num_list = precip_num_text.split()
    
    file = cfg.work_dir+'precip_hrs.txt'
    data = open(file, 'r')
    precip_hrs_txt = data.readline()
    precip_hrs = precip_hrs_txt.split()


#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# DAY 4 TEMPORARILY COMMENTED OUT !
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! 
    try:
        precip_num_list_d4 = cfg.precip_num_list_d4
        precip_d4_hrs = cfg.precip_d4_hrs
        rdict = cfg.rdict
        rdict_d4 = cfg.rdict_d4
    except:
        print(traceback.format_exc())
        print ("DAY 4 ERROR")
        logging.error(traceback.format_exc())
        logging.error("DAY 4 ERROR")
        pass

    all_hrs = []
    for hr in precip_hrs:
        all_hrs.append(hr)
    for d4hr in precip_d4_hrs:
        all_hrs.append(d4hr)
    
    rlist_6hr = gfunc.rolling_periods(fp = precip_hrs[-1], rHR=4, step=6, hr_list = precip_hrs)
    rlist_nums = gfunc.rolling_periods(fp = precip_num_list[-1], rHR=4, step=6, hr_list = precip_num_list)

    #--------------#
    # Arc Settings #
    #--------------#                      
    #Get Spatial Analyst License
    try:
        if arcpy.CheckExtension("Spatial") == "Available":
            arcpy.CheckOutExtension("Spatial")
        else:
            # Raise a custom exception
            raise LicenseError

    except LicenseError:
        print "Spatial Analyst license is unavailable"  

    arcpy.env.overwriteOutput = True
    arcpy.SetLogHistory(False)

    start_time = strftime("%a, %d %b %Y %H:%M:%S", gmtime())
    print "WSSI Snow Amount Start Time "+start_time
    print'Starting Snow Amount Index'

    logging.debug ('Starting Snow Amount Index')
    #-------------#
    # Divide Loop #
    #-------------#

    ##############
    # DAYs 1 - 3 #
    ##############

    for n in precip_num_list:
        sn_rate = cfg.work_dir+"sn_rate_f"+rdict[str(n)]+"hr.tif"
        arcpy.gp.Divide_sa(cfg.grb_dir+"snow"+str(n)+".flt", "6", sn_rate)
        #print "Snow Rate Generated for f"+str(n)
        logging.debug ("Snow Rate Generated for f"+str(n))

    #########
    # DAY 4 #
    #########

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# DAY 4 TEMPORARILY COMMENTED OUT !
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!  
    try:
        d4_sn_rates = []  
        for n in precip_num_list_d4:
            sn_rate = cfg.work_dir+"sn_rate_f"+rdict_d4[str(n)]+"hr.tif"
            arcpy.gp.Divide_sa(cfg.grb_dir+"wpc_d4snow"+str(n)+".flt", "6", sn_rate)
            d4_sn_rates.append(sn_rate)
            print "Snow Rate Generated for f"+rdict_d4[str(n)]
            logging.debug ("Snow Rate Generated for f"+rdict_d4[str(n)])
        d4_sn_rt_join = ';'.join(d4_sn_rates)
        d4_rmax_sn_rate = cfg.work_dir+'r24sn_rate_f96hr.tif' 
        arcpy.gp.CellStatistics_sa(d4_sn_rt_join, d4_rmax_sn_rate, "MAXIMUM", "DATA")
        print "created: "+d4_rmax_sn_rate
        logging.debug ("created: "+d4_rmax_sn_rate)
    except:
        print(traceback.format_exc())
        print ("DAY 4 ERROR")
        logging.error(traceback.format_exc())
        logging.error("DAY 4 ERROR")
        pass    
    #--------------------------------------#
    # 24 HR Rolling Average | 6 HR cadence #
    #--------------------------------------#
    
    for r in range (0,len(rlist_nums)):
        rmax = []
        for s in range (0,4):
            rmax.append(cfg.work_dir+"sn_rate_f"+rdict[str(rlist_nums[r][s])]+"hr.tif")     
        rmax_join = ';'.join(rmax)
        #print (rmax_join)
        rmax_sn_rate = cfg.work_dir+'rmax_sn_rate_f'+rdict[str(rlist_nums[r][-1])]+"hr.tif"  
        #print (rmax_sn_rate)
        arcpy.gp.CellStatistics_sa(rmax_join, rmax_sn_rate, "MAXIMUM", "DATA")

    for r in precip_hrs[3:]:
        rmax_sn_rate = cfg.work_dir+'rmax_sn_rate_f'+r+'hr.tif'
        rmax_sn_rate_fctr = cfg.in_mem+"sn_rate_fctr_f"+r
        rmax_sn_rate_fctr_rc = cfg.in_mem+"sn_rate_fctr_rc_f"+r
        sn_rate = cfg.work_dir+"r24sn_rate_f"+r+"hr.tif"
        arcpy.gp.Times_sa(rmax_sn_rate, cfg.rate_factor, rmax_sn_rate_fctr)
        arcpy.gp.Reclassify_sa(rmax_sn_rate_fctr, "Value", "0 0.01 0;0.01 0.10000000000000001 1;0.10000000000000001 0.20000000000000001 2;0.20000000000000001 0.33000000000000002 3;0.33000000000000002 0.5 4;0.5 0.75 5;0.75 1 6;1 1.25 7;1.25 1.5 8;1.5 2 9;2 10 10", rmax_sn_rate_fctr_rc, "DATA")
        arcpy.gp.Times_sa(rmax_sn_rate_fctr_rc, cfg.urbanfactor, sn_rate)
        print ('Final Sn Rate Generated for sn_rate_f'+r)
        logging.debug ('Final Sn Rate Generated for sn_rate_f'+r)
        arcpy.Delete_management(cfg.in_mem+'rmax_sn_rate_f'+r)
        arcpy.Delete_management(rmax_sn_rate)
        arcpy.Delete_management(rmax_sn_rate_fctr)
        arcpy.Delete_management(rmax_sn_rate_fctr_rc)
        
    #-----------------------------#
    # Max SN Rate for Full Peroid #
    #-----------------------------#

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# DAY 4 TEMPORARILY COMMENTED OUT !
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    tot_sn_rate_list = []
    tot_sn_rate_d4_list = []
    for s in precip_hrs:
        tot_sn_rate_list.append(cfg.work_dir+"sn_rate_f"+s+"hr.tif")
    tot_sn_rate = ';'.join(tot_sn_rate_list)
    for p in precip_d4_hrs:
        try:
            tot_sn_rate_d4_list.append(cfg.work_dir+"sn_rate_f"+p+"hr.tif")
        except:
            pass
    tot_sn_rate_d4 = ';'.join(tot_sn_rate_d4_list)
    print tot_sn_rate
    print tot_sn_rate_d4
    arcpy.gp.CellStatistics_sa(tot_sn_rate, cfg.in_mem+"total_sn_rate_max", "MAXIMUM", "DATA")
    try:
        arcpy.gp.CellStatistics_sa(tot_sn_rate_d4, cfg.in_mem+"total_sn_rate_d4_max", "MAXIMUM", "DATA")
    except:
        pass
    arcpy.gp.Times_sa(cfg.in_mem+"total_sn_rate_max", cfg.rate_factor, cfg.in_mem+"sn_rate_fctr")
    try:
        arcpy.gp.Times_sa(cfg.in_mem+"total_sn_rate_d4_max", cfg.rate_factor, cfg.in_mem+"sn_rate_fctr_d4")
    except:
        pass
    arcpy.gp.Reclassify_sa(cfg.in_mem+"sn_rate_fctr", "Value", "0 0.01 0;0.01 0.10000000000000001 1;0.10000000000000001 0.20000000000000001 2;0.20000000000000001 0.33000000000000002 3;0.33000000000000002 0.5 4;0.5 0.75 5;0.75 1 6;1 1.25 7;1.25 1.5 8;1.5 2 9;2 10 10", cfg.in_mem+"totalrate_re", "DATA")
    try:
        arcpy.gp.Reclassify_sa(cfg.in_mem+"sn_rate_fctr_d4", "Value", "0 0.01 0;0.01 0.10000000000000001 1;0.10000000000000001 0.20000000000000001 2;0.20000000000000001 0.33000000000000002 3;0.33000000000000002 0.5 4;0.5 0.75 5;0.75 1 6;1 1.25 7;1.25 1.5 8;1.5 2 9;2 10 10", cfg.in_mem+"totalrate_d4_re", "DATA")
    except:
        pass
    arcpy.gp.Times_sa(cfg.in_mem+"totalrate_re", cfg.urbanfactor, cfg.work_dir+"overall_max_sn_rate.tif")
    try:
        arcpy.gp.Times_sa(cfg.in_mem+"totalrate_d4_re", cfg.urbanfactor, cfg.work_dir+"overall_thru_d4_max_sn_rate.tif")
    except:
        pass
    arcpy.Delete_management(cfg.in_mem+"total_sn_rate_max")
    arcpy.Delete_management(cfg.in_mem+"sn_rate_fctr")
    arcpy.Delete_management(cfg.in_mem+"totalrate_re")

    #------------------------#
    # Snow Amount Processing #
    #------------------------#
    
    fcst_hr_list = precip_hrs[3:]

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# DAY 4 TEMPORARILY COMMENTED OUT !
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    if path.exists(cfg.grb_dir+"ds.wpc_d4snow.bin") == True:
        fcst_hr_list.append('96')
    for r in fcst_hr_list:
        #----------------------------#
        # Snow Climo Processing Loop #
        #----------------------------#
        sn_gt0 = cfg.in_mem+"sn_total_gt_0_f"+r
        sn_mask = cfg.in_mem+"sn_total_mask_f"+r
        r24sn = cfg.work_dir+"snow_r24h_f"+r+'.tif'
        r24sn_f = cfg.in_mem+"rsn_f_f"+r
        zval_top = cfg.in_mem+"zval_t_f"+r
        zvals = cfg.in_mem+"zvals_f"+r
        zvals_re = cfg.in_mem+"zvals_re_f"+r
        sn_amt_urb = cfg.work_dir+"sn_amt_f"+r+".tif"
        arcpy.gp.GreaterThan_sa(r24sn, "0", sn_gt0)
        arcpy.gp.ExtractByAttributes_sa(sn_gt0, "\"VALUE\" = 1", sn_mask)
        arcpy.Delete_management(sn_gt0)
        arcpy.gp.ExtractByMask_sa(r24sn, sn_mask, r24sn_f)
        arcpy.Delete_management(sn_mask)
        arcpy.gp.Minus_sa(r24sn_f, cfg.wssi_2dsn_avg, zval_top)
        arcpy.Delete_management(r24sn_f)
        arcpy.gp.Divide_sa(zval_top, cfg.wssi_2dsn_std, zvals)
        arcpy.Delete_management(zval_top)
        arcpy.gp.Reclassify_sa(zvals, "Value", "-1000 -100.0 1;-100.0 -1.1 2;-1.1 -0.45 3;-0.45 0.2 4;0.2 1.35 5;1.35 2.5 6;2.5 4 7;4 5.5 8;5.5 10 9;10 1000 10", zvals_re, "NODATA")
        arcpy.Delete_management(zvals)
        arcpy.gp.Times_sa(zvals_re, cfg.urbanfactor, sn_amt_urb)
        arcpy.Delete_management(zvals_re)
        arcpy.gp.CellStatistics_sa(cfg.work_dir+"r24sn_rate_f"+r+"hr.tif;"+cfg.work_dir+"sn_amt_f"+r+".tif;", cfg.work_dir+"v"+r+"hrsnamt_I.tif", "MAXIMUM", "DATA")
        print ("rolling 24 hour snow amount for "+cfg.work_dir+"v"+r+"hrsnamt_I.tif")
        logging.debug ("rolling 24 hour snow amount for "+cfg.work_dir+"v"+r+"hrsnamt_I.tif")
        # full period total snow

    ##############
    # DAYs 1 - 3 #
    ##############    

    try:
        arcpy.gp.GreaterThan_sa(cfg.work_dir+"total_snow.tif", "0", cfg.in_mem+"sn_total_gt_0")
        arcpy.gp.ExtractByAttributes_sa(cfg.in_mem+"sn_total_gt_0", "\"VALUE\" = 1", cfg.in_mem+"sn_total_mask")
        arcpy.gp.ExtractByMask_sa(cfg.work_dir+"total_snow.tif", cfg.in_mem+"sn_total_mask", cfg.in_mem+"sn_total_f")
        arcpy.gp.Minus_sa(cfg.in_mem+"sn_total_f", cfg.wssi_2dsn_avg, cfg.in_mem+"zval_total_t")
        arcpy.gp.Divide_sa(cfg.in_mem+"zval_total_t", cfg.wssi_2dsn_std, cfg.work_dir+"zval_total.tif")
        arcpy.gp.Reclassify_sa(cfg.work_dir+"zval_total.tif", "Value", "-1000 -100.0 1;-100.0 -1.1 2;-1.1 -0.45 3;-0.45 0.2 4;0.2 1.35 5;1.35 2.5 6;2.5 4 7;4 5.5 8;5.5 10 9;10 1000 10", cfg.work_dir+"total_snow_re.tif", "NODATA")
        arcpy.Delete_management(cfg.in_mem+"sn_total_gt_0")
        arcpy.Delete_management(cfg.in_mem+"sn_total_mask")
        arcpy.Delete_management(cfg.in_mem+"sn_total_f")
        arcpy.Delete_management(cfg.in_mem+"zval_total_t")
    except Exception as e:
        #print ("error in new snow climo processing total_snow_re")
        print(traceback.format_exc())
        logging.error(traceback.format_exc())

    arcpy.gp.Times_sa(cfg.work_dir+"total_snow_re.tif", cfg.urbanfactor, cfg.work_dir+"total_snow_urb.tif")
    arcpy.gp.CellStatistics_sa(cfg.work_dir+"overall_max_sn_rate.tif;"+cfg.work_dir+"total_snow_urb.tif", cfg.work_dir+"total_snamt_I.tif", "MAXIMUM", "DATA")
    #print ("overall snow amount finished ")

    
##    ##########
##    # DAYs 4 #
##    ########## 
##
###!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
### DAY 4 TEMPORARILY COMMENTED OUT !
###!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
##
##    try: 
##        arcpy.gp.GreaterThan_sa(cfg.work_dir+"total_snow_thru_d4.tif", "0", cfg.in_mem+"d4_sn_total_gt_0")
##        arcpy.gp.ExtractByAttributes_sa(cfg.in_mem+"d4_sn_total_gt_0", "\"VALUE\" = 1", cfg.in_mem+"d4_sn_total_mask")
##        arcpy.gp.ExtractByMask_sa(cfg.work_dir+"total_snow_thru_d4.tif", cfg.in_mem+"d4_sn_total_mask", cfg.in_mem+"d4_sn_total_f")
##        arcpy.gp.Minus_sa(cfg.in_mem+"d4_sn_total_f", cfg.wssi_2dsn_avg, cfg.in_mem+"d4_zval_total_t")
##        arcpy.gp.Divide_sa(cfg.in_mem+"d4_zval_total_t", cfg.wssi_2dsn_std, cfg.work_dir+"d4_zval_total.tif")
##        arcpy.gp.Reclassify_sa(cfg.work_dir+"d4_zval_total.tif", "Value", "-1000 -100.0 1;-100.0 -1.1 2;-1.1 -0.45 3;-0.45 0.2 4;0.2 1.35 5;1.35 2.5 6;2.5 4 7;4 5.5 8;5.5 10 9;10 1000 10", cfg.work_dir+"d4_total_snow_re.tif", "NODATA")
##        arcpy.Delete_management(cfg.in_mem+"d4_sn_total_gt_0")
##        arcpy.Delete_management(cfg.in_mem+"d4_sn_total_mask")
##        arcpy.Delete_management(cfg.in_mem+"d4_sn_total_f")
##        arcpy.Delete_management(cfg.in_mem+"d4_zval_total_t")
##    except Exception as e:
##        print ("error in new snow climo processing total_snow_re")
##        print(traceback.format_exc())
##        logging.error(traceback.format_exc())
##
##    arcpy.gp.Times_sa(cfg.work_dir+"d4_total_snow_re.tif", cfg.urbanfactor, cfg.work_dir+"d4_total_snow_urb.tif")
##    arcpy.gp.CellStatistics_sa(cfg.work_dir+"overall_thru_d4_max_sn_rate.tif;"+cfg.work_dir+"d4_total_snow_urb.tif", cfg.work_dir+"total_thru_d4_snamt_I.tif", "MAXIMUM", "DATA")
    print ("overall snow amount finished ")
    logging.debug ("overall snow amount finished ")


    #----------------------------------------------------------------------------------------------------#
    # Copy total, day1 and day2 output to components output folder. Day3 will be coppied in summary step #
    #----------------------------------------------------------------------------------------------------#
    os.chdir(cfg.work_dir)
    if precip_num_list[-1] == str(11):
        d3 = "v66hrsnamt_I.tif"

    else:
        d3 = "v72hrsnamt_I.tif"
    infiles = ["total_snamt_I.tif","v24hrsnamt_I.tif","v48hrsnamt_I.tif",d3,"v96hrsnamt_I.tif"]
    outfiles = ["snow_fp.tif","snow_d1.tif","snow_d2.tif","snow_d3.tif","snow_d4.tif"]
    for i,o in zip(infiles, outfiles):
        try:
            copyfile(cfg.work_dir+i, cfg.comp_out+o)
        except:
            logging.error(traceback.format_exc())
            pass            

    rolling_web_list = []
    for file in glob.glob("*hrsnamt_I.tif"):
        rolling_web_list.append(file)
    print rolling_web_list

    for f in rolling_web_list:
        rf = 'r24h_SnowAmount_f'+f[1:3]+'.tif'
        copyfile(cfg.work_dir+f, cfg.roll_out+rf)
    
    logging.debug ('Finished Snow Amount Index')


    arcpy.CheckInExtension('Spatial')
    sn_amt_end = strftime("%a, %d %b %Y %H:%M:%S", gmtime())
    print ("WSSI Snow Amount Start Time "+sn_amt_start)
    print ("WSSI Snow Amount End Time "+sn_amt_end)
    logging.debug ("Snow Amount Processing Start Time "+start_time)
    logging.debug ("Snow Amount Processing Script Finished "+sn_amt_end)
    
##############################################################################################################################################              
# Execute Function #
snow_amount()
logging.shutdown()