#! /usr/bin/env python
"""

"""

# region // imports
import sys
import os

from SDToolBox import outputmessages as outputmessages
from SDToolBox import data_acquisition
from netCDF4 import Dataset
import numpy as np

# endregion

# region // variables

# endregion


class ExtractSeaLevelPressureERA5:

    __lon_key = 'longitude'
    __lat_key = 'latitude'
    __input_lon = None
    __input_lat = None
    __input_data = None

    def __init__(self, input_data: data_acquisition.InputData) :
        """Initialize the waves extraction.
        
        Arguments:
            input_data {data_acquisition.InputData} -- Required.
        """
        # verify input_data not none
        if not input_data or \
            not input_data.coord_list or \
                len(input_data.coord_list) < 1:
            raise IOError('No valid input data.')

        input_data.lon_key = self.__lon_key
        input_data.lat_key = self.__lat_key
        self.__input_data = input_data
        self.__input_lon, self.__input_lat = input_data.coord_list[0]


    def subset_slp_era5(self, directory_path: str, year_from: int, year_to: int) :
        """Extracts a collection of netCDF4 subsets based on the input data
        set when creating the extract_sea_level_pressure object.

        Arguments:
            directory_path {str} -- Location of all the variable diretories.
            area_latitude {list} -- List with x,y coordinate for the latitude
            area_longitude {list} -- List with x,y coordinate for the longitude
            year_from {int} -- Start of time data to substract.
            year_to {int} -- End of time data to substract.

        we suppose that the input is already a point of coordinates lon and lat, since the point is
        already calculated previously (see Slp_Grd_ERA5_extractR.m lines 11-41)

        Returns:
            list(Dataset) -- collection of netCDF4 subsets per variable.
        """

        variable_name = 'msl'

        years = generate_years_array_slp(year_from, year_to)
        case_dataset_list = []
         #extracts for all the years
        for year_idx, year in enumerate(years):
            base_file_name = 'era5_Global_{}_p_{}.nc'.format(variable_name, year)
            case_file_path = os.path.join(directory_path, base_file_name)
               
                # If file does not exist simply go to the next one
            if not os.path.exists(case_file_path):
                print('File {} does not exist or could not be found.'.format(case_file_path))
                continue

            with Dataset(case_file_path, 'r', format='netCDF4') as case_dataset:
                # Get dataset
                # Find nearest point (considering we are only selecting a point)

                lat_list = case_dataset.variables[self.__lat_key][:]
                lon_list = case_dataset.variables[self.__lon_key][:]
                corr_lon = data_acquisition.get_nearest_neighbor(self.__input_lon, lon_list)
                corr_lat = data_acquisition.get_nearest_neighbor(self.__input_lat, lat_list)
                # Get variable subset
                variable_subset = case_dataset.variables[variable_name][:corr_lon:corr_lat]
                case_dataset_list.append(variable_subset)

        return case_dataset_list
        
# this method could be moved in a generic utils stript outside here
  
def generate_years_array_slp(year_from, year_to):
    years = []
    for i in range(year_to - year_from):
        years.append(year_from + i) # fills an array of years
    years.append(year_to)
    return years