# Converted from 01_planet_download.ipynb
# Load packages and connect to SentinelHub
import os
import json
import datetime
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from osgeo import gdal
from sentinelhub import MimeType, CRS, BBox, SentinelHubRequest, SentinelHubDownloadClient, \
    DataCollection, bbox_to_dimensions, DownloadRequest, SHConfig, BBoxSplitter, read_data, Geometry, SentinelHubCatalog
config = SHConfig()
catalog = SentinelHubCatalog(config=config)
import time
import shutil
import geopandas as gpd
from shapely.geometry import MultiLineString, MultiPolygon, Polygon, box, shape

# SentinelHub credentials
config.sh_client_id = '1a72d811-4f0e-4447-8282-df09608cff44'
config.sh_client_secret = 'FcBlRL29i9ZmTzhmKTv1etSMFs5PxSos'

# DataCollection setup
collection_id = 'c691479f-358c-46b1-b0f0-e12b70a9856c'
byoc = DataCollection.define_byoc(
    collection_id,
    name='planet_data2',
    is_timeless=True)

# Set some variables
project = 'kibos' # or xinavane or chemba_test_8b
# Adjust the number of days needed
days = 9 # change back to 28 which is the default. 3 years is 1095 days.
#delete all the satellite outputs -> then True
empty_folder_question = True

# Paths and folders
BASE_PATH = Path('../laravel_app/storage/app') / os.getenv('PROJECT_DIR', project)
BASE_PATH_SINGLE_IMAGES = Path(BASE_PATH / 'single_images')
folder_for_merged_tifs = str(BASE_PATH / 'merged_tif')
folder_for_virtual_raster = str(BASE_PATH / 'merged_virtual')
geojson_file = Path(BASE_PATH / 'Data' / 'pivot.geojson')

# Check if the folders exist, and if not, create them
if not os.path.exists(BASE_PATH_SINGLE_IMAGES):
    os.makedirs(BASE_PATH_SINGLE_IMAGES)
if not os.path.exists(folder_for_merged_tifs):
    os.makedirs(folder_for_merged_tifs)
if not os.path.exists(folder_for_virtual_raster):
    os.makedirs(folder_for_virtual_raster)

# Evalscript
evalscript_original = """
    //VERSION=3
    function setup() {
        return {
            input: [{
                bands: ["red", "green", "blue", "nir", "udm1"]
            }],
            output: {
                bands: 4
            }
        };
    }
    function evaluatePixel(sample) {
        var scaledBlue = [2.5 * sample.blue / 10000];
        var scaledGreen = [2.5 * sample.green / 10000];
        var scaledRed = [2.5 * sample.red / 10000];
        var scaledNIR = [2.5 * sample.nir / 10000];
        if (sample.udm1 == 0) {
            return [scaledRed, scaledGreen, scaledBlue, scaledNIR];
        } else {
            return [NaN, NaN, NaN, NaN];
        }
    }
"""

def get_true_color_request_day(time_interval, bbox, size):
    return SentinelHubRequest(
        evalscript=evalscript_original,
        input_data=[
            SentinelHubRequest.input_data(
                data_collection=DataCollection.planet_data2,
                time_interval=(time_interval, time_interval)
            )
        ],
        responses=[
            SentinelHubRequest.output_response('default', MimeType.TIFF)
        ],
        bbox=bbox,
        size=size,
        config=config,
        data_folder=str(BASE_PATH_SINGLE_IMAGES / time_interval),
    )

def download_function(slot, bbox, size):
    list_of_requests = [get_true_color_request_day(slot, bbox, size)]
    list_of_requests = [request.download_list[0] for request in list_of_requests]
    data = SentinelHubDownloadClient(config=config).download(list_of_requests, max_threads=15)
    print(f' Image downloaded for ' + slot + ' and bbox ' + str(bbox))
    time.sleep(.1)

def merge_files(slot):
    file_list = [f"{x}/response.tiff" for x in Path(BASE_PATH_SINGLE_IMAGES / slot).iterdir()]
    folder_for_merged_tifs = str(BASE_PATH / 'merged_tif' / f"{slot}.tif")
    folder_for_virtual_raster = str(BASE_PATH / 'merged_virtual' / f"merged{slot}.vrt")
    vrt_all = gdal.BuildVRT(folder_for_virtual_raster, file_list)
    vrt_all = gdal.BuildVRT(folder_for_virtual_raster, file_list)
    gdal.Translate(folder_for_merged_tifs, folder_for_virtual_raster)

days_needed = int(os.environ.get("DAYS", days))
date_str = os.environ.get("DATE")
if date_str:
    end = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
else:
    end = datetime.date.today()
start = end - datetime.timedelta(days=days_needed - 1)
slots = [(start + datetime.timedelta(days=i)).strftime('%Y-%m-%d') for i in range(days_needed)]
print('Monthly time windows:\n')
if len(slots) > 10:
    for slot in slots[:3]:
        print(slot)
    print("...")
    for slot in slots[-3:]:
        print(slot)
else:
    for slot in slots:
        print(slot)

# Download images
geo_json = gpd.read_file(str(geojson_file))
geometries = [Geometry(geometry, crs=CRS.WGS84) for geometry in geo_json.geometry]
shapely_geometries = [geometry.geometry for geometry in geometries]
bbox_splitter = BBoxSplitter(
    shapely_geometries, CRS.WGS84, (5, 5), reduce_bbox_sizes=True
)
print("Area bounding box: {}\n".format(bbox_splitter.get_area_bbox().__repr__()))
bbox_list = bbox_splitter.get_bbox_list()
info_list = bbox_splitter.get_info_list()
geometry_list = bbox_splitter.get_geometry_list()
geometry_list[0]

def is_image_available(date):
    for bbox in bbox_list:
        search_iterator = catalog.search(
            collection=byoc,
            bbox=bbox,
            time=(date, date)
        )
        if len(list(search_iterator)) > 0:
            return True
    return False

available_slots = [slot for slot in slots if is_image_available(slot)]
comparison_slots = available_slots[:min(5, len(available_slots))]
print(available_slots)
print(f"Total slots: {len(slots)}")
print(f"Available slots: {len(available_slots)}")
print(f"Excluded slots due to empty dates: {len(slots) - len(available_slots)}")

resolution = 3
for slot in available_slots:
    for bbox in bbox_list:
        bbox = BBox(bbox=bbox, crs=CRS.WGS84)
        size = bbox_to_dimensions(bbox, resolution=resolution)
        download_function(slot, bbox, size)
for slot in available_slots:
    merge_files(slot)

# Delete intermediate files
folders_to_empty = [BASE_PATH / 'merged_virtual', BASE_PATH_SINGLE_IMAGES]
def empty_folders(folders, run=True):
    if not run:
        print("Skipping empty_folders function.")
        return
    for folder in folders:
        try:
            for filename in os.listdir(folder):
                file_path = os.path.join(folder, filename)
                try:
                    if os.path.isfile(file_path):
                        os.unlink(file_path)
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                except Exception as e:
                    print(f"Error: {e}")
            print(f"Emptied folder: {folder}")
        except OSError as e:
            print(f"Error: {e}")
empty_folders(folders_to_empty, run=empty_folder_question)