# ------------------------------------------------------------------------------
# DEM download  ----------------------------------------------------------------
# ------------------------------------------------------------------------------

"""
If run as main, the different functions get called based on the
input "TESTFILENAME" path. There are 3 relevant functions:

get_prefixes: This takes as input the name of the file to get
the DEM for, and it outputs a list of associated s3 paths providing
overlapping Copernicus DEM files.

download_coreg_from_prefix: This takes the list of Copernicus DEM
file paths and it downloads the associated DEM's before aggregating
them into a TIF file coregistered to the original input. It also
does some file creation and deletion in the process, in a folder
specified in variable TEMPPATH.

get_copernicus_dem: Calls both these functions subsequently to
create the DEM based on a given input GeoTiff.
"""
from botocore.config import Config
from botocore import UNSIGNED
from pathlib import Path
import boto3
import rasterio
from osgeo import gdal
from floodsens.logger import logger

from pyproj import Proj, transform


def _get_prefixes(raster_path):
    """
    Opens raster at provided location and extracts prefixes for S3 query.

    Parameters
    ----------
    raster_path :   string or PosixPath
                    path to raster to be opened and extracted

    Returns
    ----------
    prefixes :      list of strings
                    list of prefixes for S3 download
    """
    ds_raster = rasterio.open(str(raster_path))
    epsg = str(ds_raster.read_crs().to_epsg())
    zone_letter = 'N' if epsg[2] == '6' else 'S'

    if epsg[2] == '6':
        zone_letter = 'N'
    elif epsg[2] == '7':
        zone_letter = 'S'
    else:
        zone_letter = None
    #print('Zone Letter: ', zone_letter)
    bounds = ds_raster.bounds
    #print(bounds)
    #print(epsg)

    inProj = Proj(init='epsg:'+epsg) # not yet tested if also with utm
    outProj = Proj(init='epsg:4326')

    # inProj = Proj(f'epsg:{epsg}')
    # outProj = Proj('epsg:4326')

    lower_left = transform(inProj,outProj,bounds.left,bounds.bottom)
    upper_right = transform(inProj,outProj,bounds.right,bounds.top)

    lower_left = lower_left[::-1]
    upper_right = upper_right[::-1]

    # in case Proj doesnt do utm natively
    #lower_left = utm.to_latlon(bounds.left, bounds.bottom, int(
    #    epsg[-2:]), zone_letter=zone_letter)
    #upper_right = utm.to_latlon(bounds.right, bounds.top, int(
    #    epsg[-2:]), zone_letter=zone_letter)

    if zone_letter == 'S':
        lower_left = list(lower_left)
        upper_right = list(upper_right)
        lower_left_temp = lower_left[0]
        upper_right_temp = upper_right[0]
        lower_left[0] = -(90. - upper_right_temp)
        upper_right[0] = -(90. - lower_left_temp)

    needed_northing = [int(lower_left[0])]
    if (int(lower_left[0]) - int(upper_right[0])) != 0:
        needed_northing.append(int(upper_right[0]))
    needed_easting = [int(lower_left[1])]
    if (int(lower_left[1]) - int(upper_right[1])) != 0:
        needed_easting.append(int(upper_right[1]))

    if zone_letter == 'S':
        needed_northing = list(
            range(min(needed_northing)-2, max(needed_northing)+1))
        needed_easting = list(
            range(min(needed_easting)-3, max(needed_easting)-2))
    elif zone_letter == 'N':
        needed_northing = list(
            range(min(needed_northing), max(needed_northing)+1))
        needed_easting = list(
            range(min(needed_easting)-2, max(needed_easting)+1))
    else:
        needed_northing = list(
            range(min(needed_northing), max(needed_northing)+1)
        )
        needed_easting = list(
            range(min(needed_easting), max(needed_easting)+1)
        )

    # POTENTIAL ASSUMPTION: NOT ON EQUATOR OR GREENWICH
    if needed_northing[0] >= 0:
        n_letter = 'N'
    else:
        n_letter = 'S'
        needed_northing = [abs(nn) for nn in needed_northing]

    if needed_easting[0] >= 0:
        e_letter = 'E'
    else:
        e_letter = 'W'
        needed_easting = [abs(ne) for ne in needed_easting]

    prefixes = []
    for north in needed_northing:
        north_string = str(north).zfill(2)
        for east in needed_easting:
            east_string = str(east).zfill(3)
            pfx = 'Copernicus_DSM_COG_10_' + n_letter + north_string + \
                '_00_' + e_letter + east_string + '_00_DEM'
            prefixes.append(pfx)

    return prefixes

def _download_coreg_from_prefix(prefixes,
                                out_dir='.',
                                coregfile=None,
                                sampling_alg=gdal.gdalconst.GRA_NearestNeighbour):
    
    out_dir = Path(out_dir)
    temp_dir = Path(out_dir)/"temp"
    temp_dir.mkdir(parents=True, exist_ok=True)

    fn_dem_combo_vrt = temp_dir/'combo.vrt'

    pfx = prefixes[0]
    fn_dem_utm_fin = out_dir/"10_DEM.tif"

    s3_resource = boto3.resource('s3', config=Config(signature_version=UNSIGNED))
    my_bucket = s3_resource.Bucket('copernicus-dem-30m')

    downloaded = []

    for prefix in prefixes:
        try:
            object_path = str(temp_dir/f"{prefix}.tif")
            my_bucket.download_file(f"{prefix}/{prefix}.tif", object_path)
            downloaded.append(str(object_path))
        except:
            logger.exception(f"Error when downloading prefix {prefix}")

    if len(downloaded) > 1:
        gdal.BuildVRT(str(fn_dem_combo_vrt), downloaded)

    if coregfile is None:
        gdal.Translate(str(fn_dem_utm_fin), str(fn_dem_combo_vrt))
        [file_to_delete.unlink() for file_to_delete in Path(temp_dir).glob('*')]
        temp_dir.rmdir()

    if coregfile:
        raise NotImplementedError("Coregistration not correctly implemented. Skipped!")
        warp_options = gdal.WarpOptions(resampleAlg = sampling_alg)
        gdal.Warp(options = warp_options)

    return fn_dem_utm_fin

def get_copernicus_dem(raster_path, out_dir, coregister=None):
    """
    Test get_dem_prefixes and download_coreg_from_prefix here.

    Parameters
    ----------
    raster_path :   string or PosixPath 
                    Path to (preferably GeoTiff) file to specify the region
                    where to grab the Copernicus DEM.
    out_dir :       string or PosixPath
                    Path to directory to save output DEM.
    """
    prefixes = _get_prefixes(raster_path)
    if coregister is None:
        dem_path = _download_coreg_from_prefix(prefixes,
                                            out_dir=out_dir)
    
    else:
        dem_path = _download_coreg_from_prefix(prefixes,
                                            out_dir=out_dir,
                                            coregfile=raster_path,
                                            sampling_alg = gdal.gdalconst.GRA_Cubic)
    return dem_path
