import gdal
import numpy as np
import glob
import os
import argparse
import utility_functions as utility_functions

import datetime as dt



if __name__ == '__main__':
    
    # https://eroslab.cr.usgs.gov/SCIENCE/urban-heat-islands/-/blob/5589d7a9c4a830459f14a3abd3bc90fa8e365224/Phase2/uhi_stats_classbased.py
    # https://stackabuse.com/command-line-arguments-in-python/
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--iata_code", "-b", help="iata_code")
    args = parser.parse_args()
    iata_code = args.iata_code

    #indir = r'\\igskmncnfs016.cr.usgs.gov\lsrdfs1\UHI\data\PreHPCTesting\INTENSITY_ANNUAL'
    #outDir = r'\\igskmncnfs016.cr.usgs.gov\lsrdfs1\UHI\data\PreHPCTesting\INTENSITY_STACK'
    #stat = '' #MAXLST
    #iata = 'FSD'

    indir = "/caldera/projects/usgs/eros/urban_heat_islands/Phase1/"+iata_code+"/output_folder/INTENSITY_ANNUAL"
    outDir = "/caldera/projects/usgs/eros/urban_heat_islands/Phase1/"+iata_code+"/output_folder/INTENSITY_STACK"
    stat = '' #MAXLST
    iata = iata_code

    #inlist = glob.glob(indir + os.sep + stat + os.sep + "*.tif")
    inlist = glob.glob(indir + os.sep + stat + os.sep + "*MAXLST*.tif")

    intensity_array = []

    for n in range(0, len(inlist)):
        img = gdal.Open(inlist[n])
        data = img.GetRasterBand(1).ReadAsArray().astype('float16')  # (n rows by n cols array)
        data[data==-9999] = np.nan
        intensity_array.append(data)

    intensity_array = np.stack(intensity_array)

    trans, prj = utility_functions.get_geo(inlist[0])
    x_res = img.RasterXSize
    y_res = img.RasterYSize

    processing_date = dt.datetime.today().strftime('%Y%m%d')

    #INTENSITY STACK MEAN

    mean = np.nanmean(intensity_array, axis=0)

    save_file_name = 'UHI_CU_'+iata+'_STACK_' + processing_date + '_C01_MAXLST_UNDIFF_MEAN.tif'
    #save_file_name = "UHI_"+regional_grid_name+"_"+IATA_code+"_"+representative_year+"_"+processing_year+"_"+collection_number+"_"+data_product+".tif"

    save_path = os.path.join(outDir, save_file_name)
    utility_functions.saveGeoTiff(save_path,
                                mean,
                                [x_res, y_res, 1], trans, prj, dtype=gdal.GDT_Float32)

    #INTENSITY STACK MAX
    max = np.nanmax(intensity_array, axis=0)

    save_file_name = 'UHI_CU_'+iata+'_STACK_' + processing_date + '_C01_MAXLST_UNDIFF_MAX.tif'
    #save_file_name = "UHI_"+regional_grid_name+"_"+IATA_code+"_"+representative_year+"_"+processing_year+"_"+collection_number+"_"+data_product+".tif"

    save_path = os.path.join(outDir, save_file_name)
    utility_functions.saveGeoTiff(save_path,
                                max,
                                [x_res, y_res, 1], trans, prj, dtype=gdal.GDT_Float32)

    #INTENSITY STACK MEDIAN
    median = np.nanmedian(intensity_array, axis=0)

    save_file_name = 'UHI_CU_'+iata+'_STACK_' + processing_date + '_C01_MAXLST_UNDIFF_MEDIAN.tif'
    #save_file_name = "UHI_"+regional_grid_name+"_"+IATA_code+"_"+representative_year+"_"+processing_year+"_"+collection_number+"_"+data_product+".tif"

    save_path = os.path.join(outDir, save_file_name)
    utility_functions.saveGeoTiff(save_path,
                                median,
                                [x_res, y_res, 1], trans, prj, dtype=gdal.GDT_Float32)

    ABC = None
