# Run daily at 1am
# Not yet set up to just update every day but runs everything from june 1 every time


# Here are the Modules you will need
import os
import numpy as np
import pandas as pd
from datetime import datetime,timedelta
import json

import xarray as xr
import numcodecs as ncd
import boto3
from botocore import UNSIGNED
from botocore.config import Config

# Declare initial variables
region = 'conus'
field  = 'sfc'

# set domain
bot_left_lat  = 36.8
bot_left_lon  = -114.2
top_right_lat = 38.6
top_right_lon = -110.9

# This is the file that houses the latitude and longitude data for the zarr HRRR files
latLon_data = '/uufs/chpc.utah.edu/common/home/horel-group7/Pando/hrrr/HRRR_latlon.h5'
x = xr.open_dataset(latLon_data)
lat = x.latitude.data
lon = x.longitude.data

### FUNCTIONS ###
def date_range_list(start_date, end_date, time_delta):
    # Return list of datetime.date objects between start_date and end_date by hour (inclusive).
    date_list = []
    curr_date = start_date
    while curr_date <= end_date:
        date_list.append(curr_date)
        curr_date += time_delta
    return date_list

def findChunk_andLatLons(latPoint,lonPoint):
    """
    -----------------------------------------------------------------------------
     findChunk_andLatLons: Input Lat/Lon for the chunk you desire and it returns 
                           the two ID numbers needed as well as the lats and lons
                           for the entire chunk
    -----------------------------------------------------------------------------"""
    absLat = np.abs(lat - latPoint)
    absLon = np.abs(lon - lonPoint)
    c = np.maximum(absLon,absLat)
    index = np.where(c == c.min())

    x = int(index[0])
    y = int(index[1])
    
    # Divide by chunk size - ignoring remainder
    chunkX = (x//150)
    chunkY = (y//150)
    
    x0 = chunkX*150
    y0 = chunkY*150
    xRange = x0+150
    yRange = y0+150
    
    chunkLats = lat[x0:xRange,y0:yRange]
    chunkLons = lon[x0:xRange,y0:yRange]
    
    return(chunkX, chunkY, chunkLats, chunkLons)

def create_file_path(utcDT, mdlType, lvl, variable ):
    """
    create_file_path: Creates the correct AWS file path to retrieve your data. 
    For more information about creating an aws file path see ...
        utcDT: This is the datetime object for the desired date 
        mdlType: analysis or forecast (anl, fcst)
        lvl: level of the variable i.e. (surface, 500mb)
        variable: HRRR Zarr variable
    """
    stringDate    = datetime.strftime(utcDT, '%Y%m%d')
    stringDate_hr = datetime.strftime(utcDT, '%Y%m%d_%H')

    aws_fp = 'sfc/%s/%sz_%s.zarr/%s/%s/%s/%s/' %(stringDate, stringDate_hr, mdlType, lvl, variable, lvl, variable)
    return(aws_fp)

def decompressChunk(objct, surface_pres):
    """
    -----------------------------------------------------------------------------
     decompressChunk: this decompress the chunk data and returns it in a array
                      with the shape (n, 150, 150)
    -----------------------------------------------------------------------------"""
    buf = ncd.blosc.decompress(objct)

    if surface_pres == True:
        chunk = np.frombuffer(buf, dtype='<f4')
    else:
        chunk = np.frombuffer(buf, dtype='<f2')

    single_chunk = 22500
    n = int(len(chunk) / single_chunk)

    if n == 1:
        hrrr_data = np.reshape(chunk,((150,150)))
    else:
        hrrr_data = np.reshape(chunk,((n,150,150)))
    
    return(hrrr_data)

def open_aws_hrrr_zarr(s_3, file_path, sfc_prs):
    """
    This opens a hrrr zarr variable and decompresses it and returns the variable data
    s_3: the s3 instance
    file_path: the aws path to the vairable chunk 
    """
    obj = s_3.Object('hrrrzarr', file_path)

    # weird format but this gets the zarr file contents we need
    obj_info  = obj.get()['Body'].read()    

    # Decompress the hrrr data and store it in the variable
    chunk_data = decompressChunk(obj_info, sfc_prs)
    
    return(chunk_data)

def stitch_chunks(chnk1, chnk2, lts1, lts2, lns1, lns2, c_ax, mdl_typ):    
    """
    -----------------------------------------------------------------------------
     stitch_chunks: This function is ran from the order_chunks function. After It gets
                    the chunks in the correct order and stitches them together
    -----------------------------------------------------------------------------"""
    lns3  = np.concatenate((lns1,  lns2), axis=c_ax)
    lts3  = np.concatenate((lts1,  lts2), axis=c_ax)
    
    # If its the fcst run its 3-D so need to increase the axis by 1
    if mdl_typ == True:
        c_ax += 1
    chnk3 = np.concatenate((chnk1, chnk2),axis=c_ax)
    
    return(chnk3, lns3, lts3)

def order_chunks(chunk1, chunk2, lats1, lats2, lons1, lons2, bool_fcst):
    """
    -----------------------------------------------------------------------------
     order_chunks: This recieves 2 chunks and gets them in the correct order to  
                   be combined. Fcst files are handeled a little different so 
                   bool_fcst needs to be True if 
                
       IMPORTANT
                1. CHUNKS MUST BE NEXT TO EACH OTHER TO BE COMBINED
                2. NORMAL CHUNKS HAVE DIMENSION SHAPE (150, 150) IF YOU WANT TO 
                   COMBINE LARGER (ALREADY COMBINED CHUNKS) THEN THEY NEED TO HAVE
                   THE SAME SHAPE. UNLESS ITS JUST A STRAIGHT LINE OF CHUNKS YOU 
                   CAN APPEND IT ON TO THE END OR BEGINNING OF IT.
    -----------------------------------------------------------------------------"""
    dif_x = lons1[0][0] - lons2[0][0]
    dif_y = lats1[0][0] - lats2[0][0]
    #print(dif_x)
    #print(dif_y)
    
    # find which difference has a bigger value 
    if ( abs(dif_x) < abs(dif_y) ):
        # this implies that they are stacked on top of one another   
        if dif_y < 0:
            #1st is more north
            chunk3, lons3, lats3 = stitch_chunks(chunk1, chunk2, lats1, lats2, lons1, lons2, 0, bool_fcst)
        else:
            #2nd is more north
            chunk3, lons3, lats3 = stitch_chunks(chunk2, chunk1, lats2, lats1, lons2, lons1, 0, bool_fcst)
    
    # The chunks are next to each other     
    else:
        # find out which one is furthest to the east
        if dif_x > 0:
            #2nd one is more east
            chunk3, lons3, lats3 = stitch_chunks(chunk2, chunk1, lats2, lats1, lons2, lons1, 1, bool_fcst)
        else:
            #1st one is more east
            chunk3, lons3, lats3 = stitch_chunks(chunk1, chunk2, lats1, lats2, lons1, lons2, 1, bool_fcst)
            
    return(chunk3, lons3, lats3)

def combine_2_chunks(data_dict, lat_lon, bool_fcst):
    # append the two southern chunks
    data, newLons, newLats = order_chunks(data_dict['SE'], data_dict['SW'], lat_lon['Lats_SE'], lat_lon['Lats_SW'], lat_lon['Lons_SE'], lat_lon['Lons_SW'], bool_fcst)

    return(data, newLons, newLats)

def get_2_chunks(s_3, aws_File_Paths, fcstHrs):
    """ This function will grab the data for the past fcst hours currently (6,12,18)
    and return the data in a dictionary and the lat and lons associated"""
    precipData = {}
    for aws_File_Path, fcstHr in zip(aws_File_Paths, fcstHrs):
        data = {}
        # Decompress the hrrr data and store it in the variable
        for ID, loc in zip(chnkID_dict.values(), geo_loc): 
            data[loc] = open_aws_hrrr_zarr(s_3, aws_File_Path+ID, sfc_prs)
   
        precipData[fcstHr], utah_lons, utah_lats = combine_2_chunks(data, latLon_dict, bool_fcst)


    return(precipData, utah_lons, utah_lats)

def cut_data(lat,lon,bl_lat,bl_lon,tr_lat,tr_lon):
    '''    
    #Made By Brian Blaylock see script here /uufs/chpc.utah.edu/common/home/u0553130/MS/plot_hrrr/HRRR_700.py
    Cut down full netcdf data for domain for faster plotting.
    input: the bottom left corner and top right corner lat/lon coordinates
        bl_lat = bottom left latitude
        tr_lat = top right latitude
        bl_lon = bottom left longitude
        tr_lon = top right longitude
    return: the max and min of each the arrays x and y coordinates    
    
	lat and lon are global variables of the grids lat and lon
    '''
    lat_limit = np.logical_and(lat>bl_lat,lat<tr_lat)
    lon_limit = np.logical_and(lon>bl_lon,lon<tr_lon)
    
    total_limit = np.logical_and(lat_limit,lon_limit)
    
    xmin = np.min(np.where(total_limit==True)[0])-5 # +/- a buffer to cover map area (prevents white space in plot area) 
    xmax = np.max(np.where(total_limit==True)[0])+5
    ymin = np.min(np.where(total_limit==True)[1])-5
    ymax = np.max(np.where(total_limit==True)[1])+5
       
    return xmin,xmax,ymin,ymax

def convert_precip_data(accum_precip_data):
    # We are only concerned with non zero values so make all zero values equal to NaN
    # commented out because incorrect way to handle
    #accum_precip_data[accum_precip_data == 0] = np.NAN
    
    # Precip is in units kg/m^2 where 1 kg over 1m^2 is 1mm thick. Therefore units need converted between 1mm --> inches
    # https://www.researchgate.net/post/How-do-I-convert-ERA-Interim-precipitation-estimates-from-kg-m2-s-to-mm-day
    # That conversion is 25.4mm = 1 in 
    #accum_precip_data = accum_precip_data/25.4 # That conversion is 25.4mm = 1 in 
    accum_precip_data = accum_precip_data/10 # cm

    return(accum_precip_data)

#-----------------------------------------------------------------------------#
###  INITIALIZE VARIABLES

#location: lat & lon for the four chunks for Utah
geo_loc = ['SE', 'SW']
pts_Lat = [  36,   36]
pts_Lon = [-112, -117]

# Dictionary to store lat and lon data for each of the 4 chunks
latLon_dict = {}
# dictionary to store the chunk id for each of the 4 chunks
chnkID_dict = {}

anl_hours = np.array([18,19,20,21]) #(12-15 MDT)

mdl_type = 'anl'
sfc_prs = False
    
if mdl_type == 'anl': bool_fcst = False
else: bool_fcst = True

for pt_Lat, pt_Lon, loc in zip(pts_Lat,pts_Lon, geo_loc):
    # Find the chunk for the given point
    x, y, newLats, newLons = findChunk_andLatLons(pt_Lat,pt_Lon)

    # Put it in the correct format
    if mdl_type == 'anl':
        chnkID_dict['IDs_' + str(loc) ] = '{}.{}'.format(x,y)
    else:
        chnkID_dict['IDs_' + str(loc) ] = '0.{}.{}'.format(x,y)
    # get the lat and lons for this chunk and store them in the dictionary
    latLon_dict['Lats_' + loc] = newLats
    latLon_dict['Lons_' + loc] = newLons

# Use a boto3 resource to connect to AWS sapce
s3 = boto3.resource( service_name='s3', region_name='us-west-1', config=Config(signature_version=UNSIGNED) )


# Enter Code to get start and end variables
start = datetime.today().replace(month=6, day=1, hour=0)
end   = datetime.today().replace( hour=0) - timedelta(days=1) # have it finish on the day before since this script runs in the morning
#redo 2022
start = datetime.today().replace(year=2023,month=6, day=15, hour=0)
end = datetime.today().replace(year=2023,month=9, day=27, hour=0)
print(start,end)
iterDates  = date_range_list(start, end, timedelta(days=1))

avg_pwat = {}
max_cape = {}

# go through each date
for date in iterDates:
    print('working on date: ', date)
    stringDate = datetime.strftime(date, '%m/%d/%Y')
    # We want to go from midnight to midnight local time and since data is stored in UTC
    # Convert local to UTC by +6hrs, no daylight savings so we are good
    date += timedelta(hours=6)
    # empty array to store total afternoon (12-15 MDT) data in
    cape_anl_4hr = np.zeros([150,300])
    pwat_anl_4hr = np.zeros([150,300])

    # get each hour for daily sums
    for i_hr in anl_hours:
        date = date.replace(hour=i_hr)
        # Grab CAPE Data
        try: 
            awsFilePath = create_file_path(date, mdl_type, 'surface', 'CAPE')
            cape_data, utahLons, utahLats = get_2_chunks(s3, [awsFilePath], [0]) 
            cape_anl_4hr = np.fmax(cape_anl_4hr, cape_data[0]) # combine arrays and keep max values ignore nan values
        except Exception:
            print('ERROR with CAPE on hour: ', date)
        # Grab PWAT Data
        try: 
            awsFilePath = create_file_path(date, mdl_type, 'entire_atmosphere_single_layer', 'PWAT')
            pwat_data, utahLons, utahLats = get_2_chunks(s3, [awsFilePath], [0]) 
            pwat_anl_4hr += pwat_data[0]
        except Exception:
            print('ERROR with PWAT on hour: ', date)
        
    
    # cut data down to inner domain (southwestern Utah)
    xmin,xmax,ymin,ymax = cut_data(utahLats,utahLons, bot_left_lat, bot_left_lon, top_right_lat, top_right_lon)

    pwat_anl_4hr = convert_precip_data(pwat_anl_4hr[xmin:xmax,ymin:ymax])
    cape_anl_4hr = cape_anl_4hr[xmin:xmax,ymin:ymax]
    
    # PWAT: Get average per grid point per hour
    avg_pwat[stringDate] = np.nanmean(pwat_anl_4hr) / 4 # Get averaged per hour
    # CAPE: Get average per grid point of maximum values from that day
    max_cape[stringDate] = np.nanmean(cape_anl_4hr)


# Save the data 
with open("/uufs/chpc.utah.edu/common/home/u0035056/public_html/monsoon_SW_Utah/data/hrrr_PWAT_anl_daily_avg_%s.json" %( start.year), "w") as outfile:
    json.dump(avg_pwat, outfile)

# Save the data 
with open("/uufs/chpc.utah.edu/common/home/u0035056/public_html/monsoon_SW_Utah/data/hrrr_CAPE_anl_daily_max_%s.json" %( start.year), "w") as outfile:
    json.dump(max_cape, outfile)