# -*- coding: UTF-8 -*-
# Python
"""
11-09-2023
@author: jeremy auclair

Main script to run input scripts to generate input data
"""

if __name__ == '__main__':
    
    import sys, os  # system management
    currentdir = os.path.dirname(os.path.realpath(__file__))
    sys.path.insert(0, os.path.dirname(currentdir))
    
    
    from dask.distributed import Client # to parallelise calculations
    import webbrowser  # to open dask dashboard
    from modspa_pixel.config.config import config  # to import config file
    from modspa_pixel.source.code_toolbox import format_duration  # to print memory requirements
    from modspa_pixel.preprocessing.input_toolbox import prepare_directories, set_eodag_config_file, calculate_Kcb_max_obs
    from modspa_pixel.preprocessing.download_ERA5_weather import request_ER5_weather
    from time import time  # to time code excecution
    
    # Get start time of script
    t0 = time()
    
    # Declare paths
    config_file = currentdir + os.sep + 'config' + os.sep + 'config_modspa.json'
    
    # Open config file and load parameters
    config_params = config(config_file)
    
    # Import functions depending on run mode
    mode = config_params.mode
    if mode == 'pixel':
        from modspa_pixel.preprocessing.calculate_ndvi import calculate_ndvi, interpolate_ndvi
    else:
        from modspa_pixel.preprocessing.calculate_ndvi import calculate_ndvi_parcel
        from modspa_pixel.preprocessing.extract_ndvi import extract_ndvi_stats, filter_raw_ndvi, interpolate_ndvi_parcel
        from modspa_pixel.preprocessing.parcel_to_pixel import convert_dataframe_to_xarray
        
    preferred_provider = config_params.preferred_provider
    if preferred_provider == 'usgs':
        from modspa_pixel.preprocessing.download_landsat import download_landsat_data
    else:
        from modspa_pixel.preprocessing.download_S2 import download_S2_data, extract_zip_archives
        
    
    # Run parameters
    max_cpu = config_params.max_cpu
    run_name = config_params.run_name
    start_date = config_params.start_date
    end_date = config_params.end_date
    ndvi_overwrite = config_params.ndvi_overwrite
    resolution = config_params.resolution
    
    # Path parameters
    download_path = config_params.download_path
    shapefile_path = config_params.shapefile_path
    
    # Prepare directories
    prepare_directories(config_file)
    
    # Updating the eodag config file if needed
    set_eodag_config_file(config_params.path_to_eodag_config_file, download_path, preferred_provider)
    
    # Generate inputs
    
    #===== NDVI =====#
    
    if preferred_provider == 'copernicus':
        image_path = download_path + os.sep + 'SCIHUB'
        ndvi_path = image_path + os.sep + 'NDVI'
        
    elif preferred_provider == 'theia':
        image_path = download_path + os.sep + 'THEIA'
        ndvi_path = image_path + os.sep + 'NDVI'
        
    elif preferred_provider == 'usgs':
        image_path = download_path + os.sep + 'USGS'
        ndvi_path = image_path + os.sep + 'NDVI'
    
    if preferred_provider == 'usgs':
        
        # Download and extract LandSat data
        csv_extract_file = ndvi_path + os.sep + run_name + os.sep + 'extract.csv'
        extracted_images = download_landsat_data(config_params.path_to_eodag_config_file, start_date, end_date, shapefile_path, download_path, csv_extract_file, collection = 'landsat_ot_c2_l2', cloud_cover_limit = config_params.cloud_cover_limit, max_cpu = max_cpu)
    
    else:
        # Download optical images
        csv_download_file = ndvi_path + os.sep + run_name + os.sep + 'download.csv'
        # raw_images = download_S2_data(start_date, end_date, preferred_provider, csv_download_file, shapefile_path, mode = mode, cloud_cover_limit = config_params.cloud_cover_limit, minimum_overlap = config_params.minimum_overlap)
        
        # Extract Zip archives
        csv_extract_file = ndvi_path + os.sep + run_name + os.sep + 'extract.csv'
        extracted_images = extract_zip_archives(image_path, csv_download_file, preferred_provider, resolution, csv_extract_file)
    
    # Calculate and interpolate NDVI
    if mode == 'pixel':
        
        client = Client()
        client
        if config_params.open_browser:
            webbrowser.open('http://127.0.0.1:8787/status', new=2, autoraise=True)

        # Calculate NDVI
        ndvi_precube = calculate_ndvi(extracted_images, config_file)
        
        # Interpolate NDVI
        ndvi_cube = interpolate_ndvi(ndvi_precube, config_file)     
        
        client.close()
        
    else:
        
        # Check if NDVI cube already exists
        if os.path.exists(ndvi_path + os.sep + run_name + os.sep + run_name + '_NDVI_cube_' + start_date + '_' + end_date + '.nc') and not ndvi_overwrite: pass
        
        # Calculate NDVI
        csv_ndvi_path_file = ndvi_path + os.sep + run_name + os.sep + 'ndvi.csv'
        ndvi_images = calculate_ndvi_parcel(csv_extract_file, ndvi_path, csv_ndvi_path_file, shapefile_path, overwrite = ndvi_overwrite, max_cpu = max_cpu)
    
        # Extract NDVI values on shapefile features
        csv_ndvi_extract_file = ndvi_path + os.sep + run_name + os.sep + 'ndvi_extract.csv'
        raw_ndvi = extract_ndvi_stats(ndvi_images, shapefile_path, csv_ndvi_extract_file, max_cpu = max_cpu)
        
        # Filter NDVI values for each feature
        csv_ndvi_filter_file = ndvi_path + os.sep + run_name + os.sep + 'ndvi_filter.csv'
        filtered_ndvi = filter_raw_ndvi(raw_ndvi, csv_ndvi_filter_file, max_cpu = max_cpu)
        
        # Interpolate NDVI values to a daily time step
        csv_ndvi_interp_file = ndvi_path + os.sep + run_name + os.sep + 'ndvi_interp.csv'
        interpolated_ndvi = interpolate_ndvi_parcel(filtered_ndvi, csv_ndvi_interp_file, start_date, end_date, max_cpu = max_cpu)
        
        # Convert DataFrame to netCDF4 file
        ndvi_cube = ndvi_path + os.sep + run_name + os.sep + run_name + '_NDVI_cube_' + start_date + '_' + end_date + '.nc'
        convert_dataframe_to_xarray(interpolated_ndvi, ndvi_cube, variables = ['NDVI'], data_types = ['u1'])
    
    # Calculate Kcb max from NDVI cube
    calculate_Kcb_max_obs(ndvi_cube, config_params.param_csv_file, config_params.land_cover_path)
    
    #===== Weather =====#
    
    # Run weather download and formatting script
    weather = request_ER5_weather(config_file, ndvi_cube, shapefile = shapefile_path, mode = mode, raw_S2_image_ref = ndvi_path + os.sep + run_name + os.sep + run_name + '_grid_reference.tif')
    
    # Print formatted runtime
    print('\nExcecution time: ', end = '')
    format_duration(time() - t0)