# -*- coding: utf-8 -*-
"""

@author: Adrien Wehrlé, Jason E. Box, GEUS (Geological Survey of Denmark and Greenland)

Clip ArcticDEM derived slopes for a given region based on a mask. 
Slopes have been processed using SNAP slopes calculator.

Function is run in default mode at the end of the script.

WARNING: all the variables temporary stored in memory will be cleared to enable 
the deletion of temporary outputs. 

"""

import rasterio
from shapely.geometry import box
import geopandas as gpd
from rasterio.mask import mask
from rasterio.warp import calculate_default_transform, reproject, Resampling
from IPython import get_ipython
import os
from osgeo import gdal, gdalconst


def extract_arcticdem(
    adem="/path/to/slope.img",
    region="NovayaZemlya",
    regional_mask="/path/to/NovayaZemlya.tif",
    outpath="/path/to/outputs",
    verbose=True,
):

    """
    INPUTS:
        adem: path of ArcticDEM .tif file [string]
        region: region to clip [string]
        regional_mask: path of the mask associated to the selected region [.tif]
        outpath: folder where to the clipped ArcticDEM [string]
        verbose: set to True to print details about processing [boolean]

    OUTPUTS:
        {outpath}/{region}_arcticdem_slope.tif: clipped ArcticDEM (EPSG: 3413) [.tif]
        {outpath}/{region}_mask_resampled.tif: resampled mask to fit clipped
                                               ArcticDEM resolution [.tif]

    """
    
    # clear all variables to enable file deletion
    get_ipython().magic("reset -sf")

    # check if running slope or aspect
    var = adem.split("/")[-1].split(".")[0]

    if verbose:
        if var == "slope":
            print("\n")
            print("Running extract_arctidem for %s... [SLOPE]" % region)
        elif var == "aspect":
            print("\n")
            print("ASPECT: Running extract_arctidem for %s... [ASPECT]" % region)
        else:
            print("\n")
            print(
                "ERROR: Wrong ArcticDEM input file. Rename to slope.img/aspect.img or modify the code"
            )

    # initialize output name
    target_crs = "EPSG: 3413"
    target_crs_name = target_crs.split(":")[1]

    out_tif = outpath + region + "_adem_" + var + ".tif"
    out_tif_3413 = outpath + region + "_arcticdem_" + var + "_temp" + ".tif"

    # delete outputs if already exists
    if os.path.exists(out_tif):
        if verbose:
            if var == "slope":
                print("WARNING: Clipped ArcticDEM derived slopes already exist...")
                print("Deleting Clipped ArcticDEM derived slopes...")
            elif var == "aspect":
                print(
                    "WARNING: Clipped ArcticDEM derived slope aspects already exist..."
                )
                print("Deleting Clipped ArcticDEM derived slope aspects...")
        os.remove(out_tif)

    if os.path.exists(out_tif_3413):
        if verbose:
            if var == "slope":
                print(
                    "WARNING: Reprojected and clipped ArcticDEM derived slopes already exist..."
                )
                print("Deleting reprojected and clipped ArcticDEM derived slopes...")
            elif var == "aspect":
                print(
                    "WARNING: Reprojected and clipped ArcticDEM derived slope aspect already exist..."
                )
                print(
                    "Deleting reprojected and clipped ArcticDEM derived slope aspect..."
                )
        os.remove(out_tif_3413)

    data = rasterio.open(adem)
    regional_mask_path = regional_mask
    regional_mask = rasterio.open(regional_mask)

    # create the bbox with regional_mask dimensions
    lower_right_corner = regional_mask.transform * (
        regional_mask.width,
        regional_mask.height,
    )
    upper_left_corner = regional_mask.transform * (0, 0)
    bbox = box(
        upper_left_corner[0],
        lower_right_corner[1],
        lower_right_corner[0],
        upper_left_corner[1],
    )

    # insert the bbox into a GeoDataFrame
    geo = gpd.GeoDataFrame({"geometry": bbox}, index=[0], crs=data.crs)

    # coordinates of the geometry so that rasterio handles it
    def getFeatures(gdf):
        """Function to parse features from GeoDataFrame in such a manner that rasterio wants them"""
        import json

        return [json.loads(gdf.to_json())["features"][0]["geometry"]]

    if verbose:
        print("Creating mask...")
    coords = getFeatures(geo)

    # clip source image using coords
    if verbose:
        print("Clipping input...")
    out_img, out_transform = mask(dataset=data, shapes=coords, crop=True)

    # mask output image
    out_img[regional_mask == 255] = 0

    # copy the metadata
    out_meta = data.meta.copy()

    out_meta.update(
        {
            "driver": "GTiff",
            "height": out_img.shape[1],
            "width": out_img.shape[2],
            "transform": out_transform,
            "crs": regional_mask.crs,
        }
    )
    if verbose:
        print("Saving output...")
    with rasterio.open(out_tif, "w", compress="deflate", **out_meta) as dest:
        dest.write(out_img)

    if verbose:
        print("Reprojecting output...")
    dst_crs = regional_mask.crs
    with rasterio.open(out_tif) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds
        )
        kwargs = src.meta.copy()
        kwargs.update(
            {"crs": dst_crs, "transform": transform, "width": width, "height": height}
        )

        with rasterio.open(out_tif_3413, "w", compress="deflate", **kwargs) as dst:
            reproject(
                source=rasterio.band(src, 1),
                destination=rasterio.band(dst, 1),
                src_transform=src.transform,
                src_crs=src.crs,
                dst_transform=transform,
                dst_crs=dst_crs,
                resampling=Resampling.nearest,
            )

    if verbose:
        print("Resampling mask to fit output resolution...")
    # source
    src_filename = regional_mask_path
    src = gdal.Open(src_filename, gdalconst.GA_ReadOnly)
    src_proj = src.GetProjection()
    src_geotrans = src.GetGeoTransform()

    # raster to match
    match_filename = out_tif_3413
    match_ds = gdal.Open(match_filename, gdalconst.GA_ReadOnly)
    match_proj = match_ds.GetProjection()
    match_geotrans = match_ds.GetGeoTransform()
    wide = match_ds.RasterXSize
    high = match_ds.RasterYSize

    # output/destination
    dst_filename = outpath + region + "mask_resampled_temp.tif"
    dst = gdal.GetDriverByName("Gtiff").Create(
        dst_filename, wide, high, 1, gdalconst.GDT_Float32
    )
    dst.SetGeoTransform(match_geotrans)
    dst.SetProjection(match_proj)

    # run
    gdal.ReprojectImage(src, dst, src_proj, match_proj, gdalconst.GRA_NearestNeighbour)

    del dst  # Flush

    if verbose:
        print("Masking output...")

    mask = rasterio.open(outpath + region + "mask_resampled_temp.tif")
    mask_data = mask.read(1)
    output = rasterio.open(out_tif_3413)
    output_data = output.read(1)
    profile_mask = mask.profile
    profile_output = output.profile  # saving metadata
    output_data[mask.read(1) == 0] = 0
    mask_data[mask_data == 0] = 255
    profile_mask.update(nodata=255)

    if var == "slope":
        with rasterio.open(
            outpath + region + "_arcticdem_" + var + ".tif", "w", **profile_output
        ) as dst:
            dst.write(output_data, 1)
        with rasterio.open(
            outpath + region + "_mask_resampled.tif", "w", **profile_output
        ) as dst:
            dst.write(mask_data, 1)

    if var == "aspect":
        with rasterio.open(
            outpath + region + "_arcticdem_" + var + ".tif", "w", **profile_output
        ) as dst:
            dst.write(output_data, 1)

    temp = outpath + region + "mask_resampled_temp.tif"

    return out_tif, out_tif_3413, temp


import os

regions = [
    "Greenland",
    "Iceland",
    "Svalbard",
    "FransJosefLand",
    "NovayaZemlya",
    "SevernayaZemlya",
    "JanMayen",
    "NorthernArcticCanada",
    "SouthernArcticCanada",
]

inpath = "/path/to/masks/"
verbose = True

for reg in regions:

    region_path = inpath + reg + ".tif"
    out_tif, out_tif_3413, temp = extract_arcticdem(
        regional_mask=region_path, region=reg
    )
    print("Deleting temporary outputs...")
    os.remove(out_tif)
    os.remove(out_tif_3413)
    os.remove(temp)
    out_tif, out_tif_3413, temp = extract_arcticdem(
        adem="/path/to/arctic_dem/aspect.img",
        regional_mask=region_path,
        region=reg,
    )
    if verbose:
        print("Deleting temporary outputs...")
    os.remove(out_tif)
    os.remove(out_tif_3413)
    os.remove(temp)

    if verbose:
        for i in range(4):
            print("\n")
