
#Load stack images as np
#load UHI LC as np
#rasterize polygon features to determine polygon zones
#average all pixels that match target polygon and are urban and store to table
#save zonal mean value to all pixels that match target polygon and are urban

import gdal
import ogr
import os
import numpy as np
import csv
import utility_functions as utility_functions
import glob
import rasterio
from rasterio.merge import merge
import argparse
from rasterio.mask import mask
import shutil
import time
from joblib import Parallel, delayed


def boundingBoxToOffsets(bbox, geot, xsize, ysize, lc):
    lcadj = 0
    if lc == True:
        lcadj = 1

    col1 = int((bbox[0] - geot[0]) / geot[1])
    col2 = int((bbox[1] - geot[0]) / geot[1]) + 1
    row1 = int((bbox[3] - geot[3]) / geot[5])
    row2 = int((bbox[2] - geot[3]) / geot[5]) + 1

    if bbox[0] < geot[0]:
        col1 = 0
    if bbox[1] > (geot[0]+xsize*geot[1]):
        radj = geot[0]+(xsize*geot[1])
        col2 = int((radj - geot[0]) / geot[1])
    if bbox[3] > geot[3]:
        row1 = 0 + lcadj
    if bbox[2] < (geot[3]+ysize*geot[5]):
        ladj = geot[3]+(ysize*geot[5])
        row2 = int((ladj - geot[3]) / geot[5]) - lcadj

    return [row1, row2, col1, col2]


def geotFromOffsets(row_offset, col_offset, geot):
    new_geot = [
        geot[0] + (col_offset * geot[1]),
        geot[1],
        0.0,
        geot[3] + (row_offset * geot[5]),
        0.0,
        geot[5]
    ]
    return new_geot


def setFeatureStats(uhi_id, min, max, mean, median, sd, sum, count, c, nulst_mean,
                    names=["id","class","min", "max", "mean", "median", "sd", "sum", "count", "nulst_mean"]):
    featstats = {
        names[0]: uhi_id,
        names[1]: c,
        names[2]: min,
        names[3]: max,
        names[4]: mean,
        names[5]: median,
        names[6]: sd,
        names[7]: sum,
        names[8]: count,
        names[9]: nulst_mean
    }
    return featstats

def ClassBasedStats_byUrbanZone(fn_raster, fn_zones, fn_class, workdir, itype,iata):

    mem_driver = ogr.GetDriverByName("Memory")
    mem_driver_gdal = gdal.GetDriverByName("MEM")
    shp_name = "temp"

    r_ds = gdal.Open(fn_raster)
    c_ds = gdal.Open(fn_class)
    p_ds = ogr.Open(fn_zones)

    lyr = p_ds.GetLayer()
    geot = r_ds.GetGeoTransform()
    cgeot = c_ds.GetGeoTransform()
    nodata = r_ds.GetRasterBand(1).GetNoDataValue()
    trans, prj = utility_functions.get_geo(fn_raster)
    xsize = r_ds.RasterXSize
    ysize = r_ds.RasterYSize
    xcsize = c_ds.RasterXSize
    ycsize = c_ds.RasterYSize

    zstats = []
    mosaic_list = []
    mosaic_list_uz = []

    p_feat = lyr.GetNextFeature()



    #Iterate through shapefile features (i.e. 5 km buffered urban zones)
    while p_feat:
        id = p_feat.GetField("UHI_ID")
        ##print(id)

        #Create temporary raster and shp layers in memory
        if p_feat.GetGeometryRef() is not None:

            if os.path.exists(shp_name):
                mem_driver.DeleteDataSource(shp_name)
            tp_ds = mem_driver.CreateDataSource(shp_name)
            tp_lyr = tp_ds.CreateLayer('polygons', None, ogr.wkbPolygon)

            # Copy the targetted feature to the new temporary shp
            tp_lyr.CreateFeature(p_feat.Clone())

            # Get the bounding box of the target feature and convert to offsets
            offsets = boundingBoxToOffsets(p_feat.GetGeometryRef().GetEnvelope(), geot, xsize, ysize, lc=False)
            new_geot = geotFromOffsets(offsets[0], offsets[2], geot)

            c_offsets = boundingBoxToOffsets(p_feat.GetGeometryRef().GetEnvelope(), cgeot, xcsize, ycsize, lc=True)

            tr_ds = mem_driver_gdal.Create( \
                "", \
                offsets[3] - offsets[2], \
                offsets[1] - offsets[0], \
                1, \
                gdal.GDT_UInt16)

            tr_ds.SetGeoTransform(new_geot)
            burn_value = int(id[-4:]) #the last 4 digits of the UHI_ID
            gdal.RasterizeLayer(tr_ds, [1], tp_lyr, burn_values=[burn_value])
            tr_array = tr_ds.ReadAsArray()
            ##print(str(offsets[2]) + "," + str(offsets[0]) + "," + str(offsets[3]-offsets[2]) + "," + str(offsets[1]-offsets[0]))
            
            
            '''
            This is the original code, which assumes the two inputs could have diff projections
            But, what is happening is c_array is larger than r_array
            and it is causing out of bounds later because c_array indexes are used as lookups for the r_array values

            so the following code modification uses the width and height from the r_array
            so it should not be out of bounds

            for this to work the r_ds and c_ds raster should be in the same projection, which they are in this case?

            r_array = r_ds.GetRasterBand(1).ReadAsArray( \
                offsets[2], \
                offsets[0], \
                offsets[3] - offsets[2], \
                offsets[1] - offsets[0])
            c_array = c_ds.GetRasterBand(1).ReadAsArray( \
                c_offsets[2], \
                c_offsets[0], \
                c_offsets[3] - c_offsets[2], \
                c_offsets[1] - c_offsets[0])
            '''

            r_array = r_ds.GetRasterBand(1).ReadAsArray( \
                offsets[2], \
                offsets[0], \
                offsets[3] - offsets[2], \
                offsets[1] - offsets[0])
            c_array = c_ds.GetRasterBand(1).ReadAsArray( \
                c_offsets[2], \
                c_offsets[0], \
                offsets[3] - offsets[2], \
                offsets[1] - offsets[0])




            uzout = workdir + os.sep + 'UrbanZones' + os.sep
            # if not os.path.exists(uzout):
            #     os.mkdir(uzout)

            uzout_indiv = uzout + os.sep + iata + '_sep'
            # if not os.path.exists(uzout_indiv):
            #     os.mkdir(uzout_indiv)

            save_path_uz = uzout_indiv + os.sep + 'uhi_urbanzone_' + str(id) + '.tif'
            mosaic_list_uz = mosaic_list_uz + [save_path_uz]
            y_res, x_res = c_array.shape
            utility_functions.saveGeoTiff(save_path_uz, tr_array, [x_res, y_res, 1], new_geot, prj, dtype=gdal.GDT_Float32,nodata_value=0)


            if r_array is not None:
                maskarray = np.ma.MaskedArray( \
                    r_array, \
                    mask=np.logical_or(r_array == nodata, np.logical_not(tr_array)))
                ##print("maskarray: " + str(np.unique(maskarray)))

                if maskarray is not None:
                    maskarray[maskarray == -9999] = np.nan
                    ##print("nanarray: " + str(np.unique(maskarray)))
                    look = np.unique(c_array)
                    for c in np.unique(c_array):
                        if False in maskarray[np.where(c_array == c)].mask:
                            all_masked = False
                        else:
                            all_masked = True

                        if all_masked:
                            continue

                        look = maskarray[np.where(c_array == c)].mask
                        look2 = np.unique(maskarray)
                        look3 = maskarray[np.where(c_array == c)]
                        matching = maskarray[np.where(c_array == c)]
                        unmasked = matching[matching.mask == False]

                        min = np.nanmin(maskarray[np.where(c_array == c)])
                        max = np.nanmax(maskarray[np.where(c_array == c)])
                        mean = np.nanmean(maskarray[np.where(c_array == c)])

                        #median = np.nanmedian(maskarray[np.where(c_array == c)])
                        median = np.nanmedian(unmasked)
                        
                        std = np.nanstd(maskarray[np.where(c_array == c)])
                        sum = np.nansum(maskarray[np.where(c_array == c)])
                        count = maskarray[np.where(c_array == c)].count() - np.count_nonzero(np.isnan(maskarray[np.where(c_array == c)]))

                        matching = maskarray[np.where(c_array < 20)]
                        unmasked = matching[matching.mask == False]

                        if len(unmasked) == 0:
                            nulst_mean = None
                        else:
                            nulst_mean = np.nanmean(maskarray[np.where(c_array < 20)])

                        if not all_masked:
                            ABC = None

                        zstats.append(setFeatureStats( \
                             id, \
                             min, \
                             max, \
                             mean, \
                             median, \
                             std, \
                             sum, \
                             count, \
                             c, \
                             nulst_mean))

                        classarray = np.copy(maskarray)
                        classarray[np.where(c_array != c)] = 'nan'
                        classarray[classarray > -9999] = mean
                        np.nan_to_num(classarray, copy=False, nan=-9999)
                        ##print(sum)


                            
                        classout = workdir + os.sep + 'LST_CLASS_STATS_' + itype + os.sep + 'ByClass'
                        # if not os.path.exists(classout):
                        #     os.mkdir(classout)

                        save_path = classout + os.sep + os.path.basename(fn_raster)[:-4] + '_ZMEAN_' + str(id) + '_CLASS' + str(c).rjust(2, '0') + '.tif'
                        
                        y_res, x_res = classarray.shape
                        utility_functions.saveGeoTiff(save_path,classarray,[x_res, y_res, 1], new_geot, prj, dtype=gdal.GDT_Float32)
                        mosaic_list = mosaic_list + [save_path]

                else:
                    zstats.append(setFeatureStats( \
                         id, \
                         nodata, \
                         nodata, \
                         nodata, \
                         nodata, \
                         nodata, \
                         nodata, \
                         nodata, \
                         nodata))
            else:
                zstats.append(setFeatureStats( \
                     id, \
                     nodata, \
                     nodata, \
                     nodata, \
                     nodata, \
                     nodata, \
                     nodata, \
                     nodata, \
                     nodata))

            tp_ds = None
            tp_lyr = None
            tr_ds = None

            p_feat = lyr.GetNextFeature()
    
    csvout = workdir + os.sep + "LST_CLASS_STATS_" + itype + os.sep + "ByCSV"
    # if not os.path.exists(csvout):
    #     os.mkdir(csvout)

    fn_csv = csvout + os.sep + os.path.basename(fn_raster)[:-4] + '_ZMEAN' + ".csv"
    col_names = zstats[0].keys()
    with open(fn_csv, 'w', newline='') as csvfile:
        writer = csv.DictWriter(csvfile, col_names)
        writer.writeheader()
        writer.writerows(zstats)

    #merge rasters
    src_files_to_mosaic = []
    src_files_to_mosaic_uz = []
    for raster in mosaic_list:
        src = rasterio.open(raster)
        src_files_to_mosaic.append(src)

    for uzraster in mosaic_list_uz:
        src_uz = rasterio.open(uzraster)
        src_files_to_mosaic_uz.append(src_uz)

    mosaic, out_trans = merge(src_files_to_mosaic)
    out_meta = src.meta.copy()
    out_meta.update({"driver": "GTiff",
                     "height": mosaic.shape[1],
                     "width": mosaic.shape[2],
                     "transform": out_trans
                     })

    
    mosaic_uz, out_trans_uz = merge(src_files_to_mosaic_uz)
    out_meta_uz = src_uz.meta.copy()
    out_meta_uz.update({"driver": "GTiff",
                     "height": mosaic_uz.shape[1],
                     "width": mosaic_uz.shape[2],
                     "transform": out_trans_uz
                     })

    rasout = workdir + os.sep + 'LST_CLASS_STATS_' + itype + os.sep + 'ByCityARD'
    #if not os.path.exists(rasout):
    #    os.mkdir(rasout)
    
    rasout2 = workdir + os.sep + 'LST_CLASS_STATS_' + itype + os.sep + 'ByCityCenter'
    # if not os.path.exists(rasout2):
    #     os.mkdir(rasout2)

    out_fp = rasout + os.sep + os.path.basename(fn_raster)[:-4] + '_ZMEAN' + '.tif'
    with rasterio.open(out_fp, "w", **out_meta) as dest:
        dest.write(mosaic)

    if "MAXLST" in out_fp:
        stat_type = "MAXLST"
    if "MEANLST" in out_fp:
        stat_type = "MEANLST"

    out_fp_uz = uzout + os.sep + 'uhi_urbanzone_' + str(id[:3]) + '.tif'
    check = None
    with rasterio.open(out_fp_uz, "w", **out_meta_uz) as dest_uz:
        dest_uz.write(mosaic_uz)
        ABC = None
    
    
    # Create outputs based on urban centers
    for uz in mosaic_list_uz:
        # urbanzone_FSD2201
        print(mosaic_list_uz)
        zone = os.path.basename(uz)[14:-4]
        print(zone)
        mosaic_list_cc = glob.glob( classout + os.sep + "*" + stat_type + "*" + zone + "*.tif" )
        print(mosaic_list_cc)

        src_files_to_mosaic_cc = []

        for ccraster in mosaic_list_cc:
            #print(ccraster)
            src_cc = rasterio.open(ccraster)
            print(src_cc)
            src_files_to_mosaic_cc.append(src_cc)

        #print(src_files_to_mosaic_cc)

        mosaic_cc, out_trans_cc = merge(src_files_to_mosaic_cc)
        out_meta_cc = src_cc.meta.copy()
        out_meta_cc.update({"driver": "GTiff",
                            "height": mosaic_cc.shape[1],
                            "width": mosaic_cc.shape[2],
                            "transform": out_trans_cc
                            })

        out_fp_cc = rasout2 + os.sep + os.path.basename(fn_raster)[:-4] + '_ZMEAN_' + zone + '.tif'
        print("OUTNAME = " + out_fp_cc)
        with rasterio.open(out_fp_cc, "w", **out_meta_cc) as dest_uz:
            dest_uz.write(mosaic_cc)
    

    ABC = None



def uhi_stats_classbased(workdir, iata,shapefile_root,uhi_lc_root,input_type):




    itype = [input_type] #, "STACK", "TOTAL", "ANNUAL"]


    uzout = workdir + os.sep + 'UrbanZones' + os.sep
    if os.path.exists(uzout):
        shutil.rmtree(uzout)
        time.sleep(3)
    if not os.path.exists(uzout):
        os.mkdir(uzout)

    uzout_indiv = uzout + os.sep + iata + '_sep'
    if not os.path.exists(uzout_indiv):
        os.mkdir(uzout_indiv)

    classout = workdir + os.sep + 'LST_CLASS_STATS_' + itype[0] + os.sep + 'ByClass'
    if not os.path.exists(classout):
        os.mkdir(classout)

    csvout = workdir + os.sep + "LST_CLASS_STATS_" + itype[0] + os.sep + "ByCSV"
    if not os.path.exists(csvout):
        os.mkdir(csvout)

    rasout = workdir + os.sep + 'LST_CLASS_STATS_' + itype[0] + os.sep + 'ByCityARD'
    if not os.path.exists(rasout):
        os.mkdir(rasout)

    rasout2 = workdir + os.sep + 'LST_CLASS_STATS_' + itype[0] + os.sep + 'ByCityCenter'
    if not os.path.exists(rasout2):
        os.mkdir(rasout2)



    for i in itype:

        if itype == "TOTAL":
            stype = ["R2", "SLOPE"]
        else:
            stype = ''

        process_tuples = []

        rasters_match_string = workdir + os.sep + "LST_" + i + "_CLIPPED" + os.sep + "*" + stype + ".tif"
        rasters = glob.glob(rasters_match_string)
        for fn_raster in rasters:

            #identify target raster IATA code

            #iata_char_start = len(workdir) + 1 + len(i) +1 + 7
            #iata_char_end = iata_char_start + 3
            #iata = fn_raster[iata_char_start:iata_char_end]

            #locate ARD based city buffered shapefiles
            
            #fn_zones = shapefile_root + os.sep + "ard_" + iata + "_city_dissolved.shp"
            #fn_zones = workdir + os.sep + "Shapefiles" + os.sep + "ard_" + iata + "_city_dissolved.shp"

            
            fn_zones = shapefile_root + os.sep + "City_Boundaries_Dissolved" + os.sep + "ard_" + iata.lower() + "_city_dissolved.shp"

            #looking = os.path.basename(rasters[0])
            
            #year = os.path.basename(rasters[0])[11:15]
            year = os.path.basename(fn_raster)[11:15]


            #year = os.path.basename(rasters[0])[17:21]

            #determine what UHI_LC layer to use (if STACK, use most recently available layer)
            #lc_raster_match_string = uhi_lc_root + os.sep + "UHI_LC" + os.sep + "*.tif"
        
            lc_raster_match_string = uhi_lc_root + os.sep + "*.tif"
            lc_rasters = glob.glob(lc_raster_match_string)

            max_lcyr = 0
            for lc in lc_rasters:
                yr = int(os.path.basename(lc)[11:15])
                if yr > max_lcyr:
                    max_lcyr = yr
            
            #/caldera/projects/usgs/eros/urban_heat_islands/Phase1/FSD/output_folder/STACK/UHI_CU_FSD_STACK_20220503_C01_MEANLST_MIN.tif
            #so in the future we have to change STACK to a year range start_year_end_year something like that to include the years
            #but it isnt decided on what it will be yet
            if year == "STAC":
                year = str(max_lcyr)

            match_string = workdir + os.sep + "UHI_LC" + os.sep + "UHI_CU_" + iata + "_" + i + "_" + year + "*LC.tif"
            match_string1 = uhi_lc_root + os.sep + "UHI_CU_" + iata + "_" + year + "*LC.tif"
            fn_class = glob.glob(match_string1)
            fn_class = fn_class[0]

            process_tuples.append((fn_raster, fn_zones, fn_class, workdir, i,iata))

        for (fn_raster, fn_zones, fn_class, workdir, i,iata) in process_tuples:
             ClassBasedStats_byUrbanZone(fn_raster, fn_zones, fn_class, workdir, i,iata)
             ABC = None

        #Parallel(n_jobs=20)(delayed(ClassBasedStats_byUrbanZone)(fn_raster, fn_zones, fn_class, workdir, i,iata) for (fn_raster, fn_zones, fn_class, workdir, i,iata) in process_tuples)

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/
    

    # /caldera/projects/usgs/eros/urban_heat_islands/Phase1/FSD/output_folder
    # FSD
    # /caldera/projects/usgs/eros/urban_heat_islands/Phase2/uhi_metrobuffer_generation


    #for iata_code in ["MIA","ABQ","MCI","SMF"]:#in progress
    #for iata_code in ["MKE"]:#in progress

    iata_code = "CLT"
    workdir = "/caldera/projects/usgs/eros/urban_heat_islands/Phase1/" + iata_code + "/output_folder"
    metrobuffer_shapefile_root = "/caldera/projects/usgs/eros/urban_heat_islands/Phase2/uhi_metrobuffer_generation"
    uhi_lc_root = "/caldera/projects/usgs/eros/urban_heat_islands/Phase2/section_332/output/" + iata_code
    input_type = "ANNUAL"
    
    #I don't think this requires anything from this loop so it can be moved to the top
    classout_path = workdir + os.sep + 'LST_CLASS_STATS_' + input_type
    if os.path.exists(classout_path):
        shutil.rmtree(classout_path)
        time.sleep(3)
    if not os.path.exists(classout_path):
        os.mkdir(classout_path)


    uhi_stats_classbased(workdir, iata_code, metrobuffer_shapefile_root, uhi_lc_root,input_type)
    



    '''
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", "-w", help="specify workpath directory")
    parser.add_argument("--iata_location_identifier", "-i", help="iata_location_identifier")
    parser.add_argument("--metrobuffer_shapefile_root", "-m", help="metrobuffer_shapefile_root") 
    parser.add_argument("--input_type", "-t", help='list of 1-3 input types ["STACK", "TOTAL", "ANNUAL"]')
    parser.add_argument("--uhi_lc_root", "-lc", help="uhi_lc_root")
    args = parser.parse_args()
    uhi_stats_classbased(args.workdir, args.iata_location_identifier, args.metrobuffer_shapefile_root, args.uhi_lc_root,args.input_type)
    '''

    #"/caldera/projects/usgs/eros/urban_heat_islands/Phase2/section_332/output/FSD"
    #uhi_stats_classbased("/caldera/projects/usgs/eros/urban_heat_islands/Phase1/FSD/output_folder", "FSD", "/caldera/projects/usgs/eros/urban_heat_islands/Phase2/uhi_metrobuffer_generation","/caldera/projects/usgs/eros/urban_heat_islands/Phase2/section_332/output/FSD","ANNUAL")
