
import json
import math
import os
import shutil
import time

import rasterio
from rasterio.mask import mask
from rasterio.io import MemoryFile
from rasterio.merge import merge

from osgeo import osr
import pycrs
import pyproj
import numpy as np
import fiona

from joblib import Parallel, delayed

import datetime as dt
import argparse





def do_step_344(raster_step_332_1985_path,raster_step_332_2020_path,final_output_directory,full_path_to_dissolved_city_boundaries_shapefile):


    if not os.path.exists(final_output_directory):
        os.mkdir(final_output_directory)


    boundary_geometries = []
    with fiona.open(full_path_to_dissolved_city_boundaries_shapefile, "r") as shapefile:
        for feature in shapefile:
            boundary_geometries.append(feature["geometry"])

    source_raster_step_332_1985 = rasterio.open(raster_step_332_1985_path)
    source_raster_step_332_2020 = rasterio.open(raster_step_332_2020_path)


    masked_raster_A, masked_transform_raster_A = mask(source_raster_step_332_1985, boundary_geometries, crop=False, pad=False, pad_width=0.0,
                                                  all_touched=True)
    masked_raster_B, masked_transform_raster_B = mask(source_raster_step_332_2020, boundary_geometries, crop=False, pad=False, pad_width=0.0,
                                                  all_touched=True)

    designated_output_nodata_val = -9999

    new_image = np.full(masked_raster_A.shape, designated_output_nodata_val,dtype=np.int32)
    new_out_meta = source_raster_step_332_1985.meta
    new_out_meta['dtype'] = 'int32'

    for n1 in range(new_image.shape[1]):
        for n2 in range(new_image.shape[2]):
            old_val = masked_raster_A[0][n1][n2]
            new_val = masked_raster_B[0][n1][n2]
            o_val = designated_output_nodata_val
            if new_val in [21,22,23,24,25] and old_val not in [21,22,23,24,25]:
                o_val = 1
            else:
                o_val = 0
            if o_val == 0 and old_val == 2 and new_val in [21,22,23,24,25]:
                o_val = -1 #should never happen
            new_image[0][n1][n2] = o_val

    new_out_meta.update({"driver": "GTiff",
                         "height": masked_raster_A.shape[1],
                         "width": masked_raster_A.shape[2],
                         "transform": masked_transform_raster_A,
                         "nodata": designated_output_nodata_val,
                         "compress": "LZW"})

    output_path = os.path.join(final_output_directory,"new_imp_85_20.tif")

    with rasterio.open(output_path, "w", **new_out_meta) as dest:
        dest.write(new_image)
        print(output_path)




'''
full_path_to_dissolved_city_boundaries_shapefile = "D:\\Users\\rhussain\\Desktop\\PycharmProjects\\prototype_uhi_phase2\\City_Boundaries_Dissolved\\ard_fsd_city_dissolved.shp"

raster_step_332_1985_path = "Y:\\UHI\\data\\phase2_test_outputs\\full_set8\\SiouxFalls\\UHI_CU_FSD_1985_20211116_C01_LC.tif"
raster_step_332_2020_path = "Y:\\UHI\\data\\phase2_test_outputs\\full_set8\\SiouxFalls\\UHI_CU_FSD_2020_20211221_C01_LC.tif"
final_output_directory = "D:\\RezaTemp\\uhi_step_334\\"

do_step_344(raster_step_332_1985_path,raster_step_332_2020_path,final_output_directory,full_path_to_dissolved_city_boundaries_shapefile)
'''


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("--full_path_to_dissolved_city_boundaries_shapefile", "-s", help="specify full_path_to_dissolved_city_boundaries_shapefile directory")
    parser.add_argument("--raster_step_332_1985_path", "-a", help="raster_step_332_1985_path")
    parser.add_argument("--raster_step_332_2020_path", "-b", help="raster_step_332_2020_path")
    parser.add_argument("--final_output_directory", "-o", help="final_output_directory")
    args = parser.parse_args()
    do_step_344(args.raster_step_332_1985_path,args.raster_step_332_2020_path,args.final_output_directory,args.full_path_to_dissolved_city_boundaries_shapefile)
