import datetime as dt
from pathlib import Path
import cloudbutton_geospatial.s2froms3 as s2froms3
import matplotlib.pyplot as plt
import numpy as np
import rasterio
import lithops
import time
import shutil
import os

# https://pybonacci.org/2021/04/21/creando-un-mosaico-usando-s2froms3-en-python-a-partir-de-imagenes-de-sentinel-2/
# http://patrickgray.me/open-geo-tutorial/chapter_2_indices.html


def copy_band_to_runtime(band_s3_location, storage):
    band_path = band_s3_location.split('/')
    band_bucket = band_path[0]
    band_key = '/'.join(band_path[1:])
    band_local = f'/tmp/{band_path[7]}_{band_path[8]}'
    band_stream = storage.get_object(band_bucket, band_key, stream=True)
    with open(band_local, 'wb') as file:
        shutil.copyfileobj(band_stream, file)
        print(f'Finished copying {band_s3_location} to local disk')
    return band_local


def ndvi(band_4_s3_location, band_8_s3_location, storage):
    # AWS lambda has only 512 of emphemeral storage
    # band4_local = copy_band_to_runtime(band_4_s3_location, storage)
    # band8_local = copy_band_to_runtime(band_8_s3_location, storage)

    # Clean /tmp dir
    for f in os.listdir('/tmp'):
        path = os.path.join('/tmp', f)
        print(path)
        os.remove(path) if os.path.isfile(path) else None

    band_path = band_4_s3_location.split('/')
    ndvi_local = f'/tmp/{band_path[7]}_NDVI.tif'

    # generate nir and red objects as arrays in float64 format
    band4 = rasterio.open('s3://'+band_4_s3_location)  # red
    band8 = rasterio.open('s3://'+band_8_s3_location)  # nir

    profile = band4.profile
    profile.update(dtype='float64')

    with rasterio.open(ndvi_local, 'w', **profile) as dst:
        for _, window in band4.block_windows(1):
            red = band4.read(1, window=window).astype('float64')
            nir = band8.read(1, window=window).astype('float64')
            ndvi = (np.where((nir + red) == 0., 0, (nir - red) / (nir + red))).astype('float64')
            dst.write(ndvi, 1, window=window)

    with open(ndvi_local, 'rb') as ndvi_file:
        co = storage.put_cloudobject(ndvi_file, key=ndvi_local.replace('/tmp/', ''))

#     fig = plt.figure(figsize=(20, 15))
#     jpg_path = os.path.join(ndvi_local.replace('tiff', 'jpg'))
#     with rasterio.open(ndvi_local) as src:
#         arr = src.read(1)
#         plt.imsave(jpg_path, arr, cmap="RdYlGn")
#
#     with open(jpg_path, 'rb') as jpg_ndvi_file:
#         co = storage.put_cloudobject(jpg_ndvi_file, key=jpg_path.replace('/tmp/', ''))

    return co


if __name__ == "__main__":

    lat = 41.115696  # Tarragona
    lon = 1.249594  # Tarragona

    lat = 39.60595289727246  # California
    lon = -122.82804126978336  # California

    start_date = dt.date(2019, 9, 16)  # Start date to search images
    end_date = dt.date(2019, 9, 17)  # End date to search images
    what = ['B04', 'B08']  # What we want to download
    cc = 10  # Minimum cloud cover on each image, 25 is 25%
    folder = Path('./AwsData')  # Where the files will be downloaded

    print(s2froms3.point_in_tile(lon, lat))

    t0 = time.time()
    scenes = s2froms3.get_scene_list(lon=lon, lat=lat, start_date=start_date, end_date=end_date, what=what, cloud_cover_le=cc)
    print(scenes)
    #downloaded = s2froms3.download_S2(scenes=scenes, folder=folder)
    #print("Downloaded: ", downloaded)

    print("\nTotal tiles found: {} ({}s)".format(len(scenes), round(time.time()-t0, 2)))

    while True:
        if input('\nDo You Want To Continue? (y/n): ') == 'y':
            break
        else:
            exit()

    fexec = lithops.FunctionExecutor(backend='aws_ec2', storage='aws_s3', log_level='DEBUG',
                                     runtime_memory=2048, runtime='lithops-awslambda-ndvi-v38:01')
    res = fexec.map(ndvi, scenes)
    cloudobjects = fexec.get_result()
    for co in cloudobjects:
        ndvi_stream = fexec.storage.get_cloudobject(co, stream=True)
        ndvi_tif = os.path.join(folder, co.key)
        with open(ndvi_tif, 'wb') as file:
            shutil.copyfileobj(ndvi_stream, file)

    print('Finished downloading NDVI tifs')

    fig = plt.figure(figsize=(20, 15))
    jpg_path = os.path.join(ndvi_tif.replace('tif', 'jpg'))
    with rasterio.open(ndvi_tif) as src:
        arr = src.read(1)
        plt.imsave(jpg_path, arr, cmap="RdYlGn")

import datetime as dt
from pathlib import Path
import cloudbutton_geospatial.s2froms3 as s2froms3
import matplotlib.pyplot as plt
import numpy as np
import rasterio
import lithops
import time
import shutil
import os

# https://pybonacci.org/2021/04/21/creando-un-mosaico-usando-s2froms3-en-python-a-partir-de-imagenes-de-sentinel-2/
# http://patrickgray.me/open-geo-tutorial/chapter_2_indices.html


def copy_band_to_runtime(band_s3_location, storage):
    band_path = band_s3_location.split('/')
    band_bucket = band_path[0]
    band_key = '/'.join(band_path[1:])
    band_local = f'/tmp/{band_path[7]}_{band_path[8]}'
    band_stream = storage.get_object(band_bucket, band_key, stream=True)
    with open(band_local, 'wb') as file:
        shutil.copyfileobj(band_stream, file)
        print(f'Finished copying {band_s3_location} to local disk')
    return band_local


def ndvi(band_4_s3_location, band_8_s3_location, storage):
    # AWS lambda has only 512 of emphemeral storage
    # band4_local = copy_band_to_runtime(band_4_s3_location, storage)
    # band8_local = copy_band_to_runtime(band_8_s3_location, storage)

    # Clean /tmp dir
    for f in os.listdir('/tmp'):
        path = os.path.join('/tmp', f)
        print(path)
        os.remove(path) if os.path.isfile(path) else None

    band_path = band_4_s3_location.split('/')
    ndvi_local = f'/tmp/{band_path[7]}_NDVI.tif'

    # generate nir and red objects as arrays in float64 format
    band4 = rasterio.open('s3://'+band_4_s3_location)  # red
    band8 = rasterio.open('s3://'+band_8_s3_location)  # nir

    profile = band4.profile
    profile.update(dtype='float64')

    with rasterio.open(ndvi_local, 'w', **profile) as dst:
        for _, window in band4.block_windows(1):
            red = band4.read(1, window=window).astype('float64')
            nir = band8.read(1, window=window).astype('float64')
            ndvi = (np.where((nir + red) == 0., 0, (nir - red) / (nir + red))).astype('float64')
            dst.write(ndvi, 1, window=window)

    with open(ndvi_local, 'rb') as ndvi_file:
        co = storage.put_cloudobject(ndvi_file, key=ndvi_local.replace('/tmp/', ''))

#     fig = plt.figure(figsize=(20, 15))
#     jpg_path = os.path.join(ndvi_local.replace('tiff', 'jpg'))
#     with rasterio.open(ndvi_local) as src:
#         arr = src.read(1)
#         plt.imsave(jpg_path, arr, cmap="RdYlGn")
#
#     with open(jpg_path, 'rb') as jpg_ndvi_file:
#         co = storage.put_cloudobject(jpg_ndvi_file, key=jpg_path.replace('/tmp/', ''))

    return co


if __name__ == "__main__":

    lat = 41.115696  # Tarragona
    lon = 1.249594  # Tarragona

    lat = 39.60595289727246  # California
    lon = -122.82804126978336  # California

    start_date = dt.date(2019, 9, 16)  # Start date to search images
    end_date = dt.date(2019, 9, 17)  # End date to search images
    what = ['B04', 'B08']  # What we want to download
    cc = 10  # Minimum cloud cover on each image, 25 is 25%
    folder = Path('./AwsData')  # Where the files will be downloaded

    print(s2froms3.point_in_tile(lon, lat))

    t0 = time.time()
    scenes = s2froms3.get_scene_list(lon=lon, lat=lat, start_date=start_date, end_date=end_date, what=what, cloud_cover_le=cc)
    print(scenes)
    #downloaded = s2froms3.download_S2(scenes=scenes, folder=folder)
    #print("Downloaded: ", downloaded)

    print("\nTotal tiles found: {} ({}s)".format(len(scenes), round(time.time()-t0, 2)))

    while True:
        if input('\nDo You Want To Continue? (y/n): ') == 'y':
            break
        else:
            exit()

    fexec = lithops.FunctionExecutor(backend='aws_ec2', storage='aws_s3', log_level='DEBUG',
                                     runtime_memory=2048, runtime='lithops-awslambda-ndvi-v38:01')
    res = fexec.map(ndvi, scenes)
    cloudobjects = fexec.get_result()
    for co in cloudobjects:
        ndvi_stream = fexec.storage.get_cloudobject(co, stream=True)
        ndvi_tif = os.path.join(folder, co.key)
        with open(ndvi_tif, 'wb') as file:
            shutil.copyfileobj(ndvi_stream, file)

    print('Finished downloading NDVI tifs')

    fig = plt.figure(figsize=(20, 15))
    jpg_path = os.path.join(ndvi_tif.replace('tif', 'jpg'))
    with rasterio.open(ndvi_tif) as src:
        arr = src.read(1)
        plt.imsave(jpg_path, arr, cmap="RdYlGn")

# https://pybonacci.org/2021/04/21/creando-un-mosaico-usando-s2froms3-en-python-a-partir-de-imagenes-de-sentinel-2/ # http://patrickgray.me/open-geo-tutorial/chapter_2_indices.html # AWS lambda has only 512 of emphemeral storage # band4_local = copy_band_to_runtime(band_4_s3_location, storage) # band8_local = copy_band_to_runtime(band_8_s3_location, storage) # Clean /tmp dir # generate nir and red objects as arrays in float64 format # red # nir #     fig = plt.figure(figsize=(20, 15)) #     jpg_path = os.path.join(ndvi_local.replace('tiff', 'jpg')) #     with rasterio.open(ndvi_local) as src: #         arr = src.read(1) #         plt.imsave(jpg_path, arr, cmap="RdYlGn") # #     with open(jpg_path, 'rb') as jpg_ndvi_file: #         co = storage.put_cloudobject(jpg_ndvi_file, key=jpg_path.replace('/tmp/', '')) # Tarragona # Tarragona # California # California # Start date to search images # End date to search images # What we want to download # Minimum cloud cover on each image, 25 is 25% # Where the files will be downloaded #downloaded = s2froms3.download_S2(scenes=scenes, folder=folder) #print("Downloaded: ", downloaded)