'''
#################################################################################
This software prototype has been developed under the Copernicus C3S_512 contract.

Institution: BSC
Author: Jesus Peña-Izquierdo
Email: jesus.pena@bsc.es
#################################################################################
'''

import xarray as xr
from pathlib import Path
import time
import os
from C3S512.cds import downloader
import cdsapi
import logging
from src.tools import tools
t = tools()


class downloading:   
    
    def __init__(self,
                 conf
                ):
        
        self.dataset_parameters = conf['dataset_parameters']
        
        self.variable = conf['dataset_parameters']['variable']
        self.originating_centre = conf['dataset_parameters']['originating_centre']
        self.system = conf['dataset_parameters']['system']
        self.years = conf['dataset_parameters']['years']
        self.start_months = conf['dataset_parameters']['start_months']
        self.leadtime_months = conf['dataset_parameters']['leadtime_months']
        self.levels = conf['dataset_parameters']['levels']
        
        self.dataset_sf = conf['dataset_parameters']['dataset_sf']
        self.product_type_sf = conf['dataset_parameters']['product_type_sf']
        self.download_aggregation_sf = conf['download_parameters']['download_aggregation_sf']
        self.file_format_sf = conf['download_parameters']['file_format_sf']
        
        self.dataset_ref = conf['dataset_parameters']['dataset_ref']
        self.product_type_ref = conf['dataset_parameters']['product_type_ref']
        self.download_aggregation_ref = conf['download_parameters']['download_aggregation_ref']
        self.file_format_ref = conf['download_parameters']['file_format_ref']  
        
        self.variable_ref = conf['dataset_parameters']['variable_ref']
        if not self.variable_ref:
            self.variable_ref = self.variable
        
        self.force_download = conf['download_parameters']['force_download']
        self.retry_download = conf['download_parameters']['retry_download']
        self.c3s512_downloader = conf['download_parameters']['c3s512_downloader']
        self.base_dataIN_path = Path(conf['folders']['base_dataIN_path'])
        self.logs_path = Path(conf['folders']['logs_path'])
                    
        formatter = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        log_error_file = 'log_error_download_{}s{}_{}.log'.format(self.originating_centre,
                                                                  self.system,
                                                                  self.variable)
        logging.basicConfig(filename = self.logs_path.joinpath(log_error_file),
                            level = "ERROR",
                            format = formatter,
                            filemode = "w")
        
        if self.c3s512_downloader:
            self.downloader = downloader
        else:
            c = cdsapi.Client()
            self.downloader = c.retrieve
        
        
        print('\n#########################')
        print('DOWNLOADING..............')
        print('#########################\n')
        
        t.print_header(self.dataset_parameters)

        
    def down(self,
             download_form,
             download_path,
             dataset
            ):

        # Call retrieve and download the file
        if self.force_download:
            download_flag = True
        elif not os.path.exists(download_path):
            download_flag = True
        else:
            download_flag = False
            print('\nFile {} download skipped.'.format(download_path))
        
        while download_flag:
            try:
                print("\nDownload request: {}".format(download_form))
                self.downloader(dataset,
                               download_form,
                               download_path #,key = '7173:160c5dd3-ad62-40ec-8bb9-6d6b25bb755d'
                              )
                
                print('\nFile {} downloaded.'.format(download_path))
                download_flag = False
                
            except Exception as e:
                print(e)
                e = "\nDownload error: {}".format(download_form)
                logging.error(e)

                if self.retry_download:
                    download_flag = True
                else:
                    download_flag = False
            


    def seasonal_forecast(self):

        if 'pressure' in self.dataset_sf:
            if self.levels is None:
                print('WARNING: Not pressure levels selected!!!')
            levels_name = 'pressure_level'
                
        
        pathIN = self.base_dataIN_path.joinpath('seasonal',
                                                  self.dataset_sf)

        if self.file_format_sf == 'grib':
            file_format_ext = 'grib'
        if self.file_format_sf == 'netcdf':
            file_format_ext = 'nc'
        
        tic = time.time()
        
        if self.download_aggregation_sf == 'monthly': #Monthly files
            
            pathIN = pathIN.joinpath('monthly_files',
                                     self.variable
                                    )
            
            if not os.path.exists(pathIN):
                os.makedirs(pathIN)
        
            for yy in range(self.years[0], self.years[1]+1):
                for mm in range(self.start_months[0], self.start_months[1]+1):
                    
                    download_file = '{}s{}-{}-{}{:02d}01.{}'.format(self.originating_centre,
                                                                 self.system,
                                                                 self.variable,
                                                                 yy, mm, 
                                                                 file_format_ext)
                    download_path = pathIN.joinpath(download_file)
                    download_form ={'variable': self.variable,
                                    'originating_centre': self.originating_centre,
                                    'system': self.system,
                                    'product_type': self.product_type_sf,
                                    'year': [str(yy)],
                                    'month': ['{:02d}'.format(mm)],
                                    'leadtime_month': self.leadtime_months,
                                    'format': self.file_format_sf
                                   }

                    if 'pressure' in self.dataset_sf:
                        download_form[levels_name] = self.levels

                    #Download request
                    self.down(download_form,
                              download_path,
                              dataset = self.dataset_sf
                             )
                    
        else: #One single file per start month
            
            pathIN = pathIN.joinpath('startmonth_files')
            
            if not os.path.exists(pathIN):
                os.makedirs(pathIN)
            
            year_limits = '{}to{}'.format(self.years[0], self.years[1])
            for mm in range(self.start_months[0], self.start_months[1] + 1):
                download_file = '{}s{}-{}-{}-s{:02d}.nc'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              year_limits,
                                                              mm
                                                             )
                download_path = pathIN.joinpath(download_file)
                download_form ={'variable': self.variable,
                                'originating_centre': self.originating_centre,
                                'system': self.system,
                                'product_type': self.product_type_sf,
                                'year': ['{}'.format(yy) for yy in range(self.years[0], self.years[1]+1)],
                                'month': ['{:02d}'.format(mm)],
                                'leadtime_month': self.leadtime_months,
                                'format': self.file_format_sf
                               }

                if 'pressure' in self.dataset_sf:
                    download_form[levels_name] = self.levels

                #Download request
                self.down(download_form,
                          download_path,
                          dataset = self.dataset_sf
                         )
            
            
    def era5(self):

        if 'pressure' in self.dataset_ref:
            if self.levels is None:
                print('WARNING: Not pressure levels selected!!!')
            levels_name = 'pressure_level'
                
        
        pathIN = self.base_dataIN_path.joinpath('reanalysis', 
                                                  self.dataset_ref
                                               )

        if self.file_format_ref == 'grib':
            file_format_ext = 'grib'
        if self.file_format_ref == 'netcdf':
            file_format_ext = 'nc'
        
        tic = time.time()
        
        if self.download_aggregation_ref == 'monthly': #Monthly files
            
            pathIN = pathIN.joinpath('monthly_files',
                                     self.variable
                                    )
            if not os.path.exists(pathIN):
                os.makedirs(pathIN)
        
            for yy in range(self.years[0], self.years[1] + 2): # Add one more year for validating the forecast
                for mm in range(self.start_months[0], self.start_months[1] + 1): 
                    
                    download_file = '{}-{}{:02d}01.{}'.format(self.variable, yy, mm, file_format_ext)
                    download_path = pathIN.joinpath(download_file)
                    if 'monthly' in self.dataset_ref:
                        download_form ={'product_type': self.product_type_ref,
                                    'variable': self.variable,
                                    'year': [str(yy)],
                                    'month': ['{:02d}'.format(mm)],
                                    'time': '00:00',
                                    'format': self.file_format_ref
                                   }
                    else:
                         download_form ={'product_type': self.product_type_ref,
                                    'variable': self.variable,
                                    'year': [str(yy)],
                                    'month': ['{:02d}'.format(mm)],
                                    'day': [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31],
                                    'time': ['00:00', '01:00', '02:00', '03:00', '04:00', '05:00', '06:00', '07:00', '08:00', '09:00', '10:00', '11:00', '12:00', '13:00', '14:00', '15:00', '16:00', '17:00', '18:00', '19:00', '20:00', '21:00', '22:00', '23:00'],
                                    'format': self.file_format_ref
                                   }
                    if 'pressure' in self.dataset_ref:
                        download_form[levels_name] = self.levels

                    #Download request
                    self.down(download_form,
                              download_path,
                              dataset = self.dataset_ref
                             )
        else: #One single file
            
            pathIN = pathIN.joinpath('single_file')
            if not os.path.exists(pathIN):
                os.makedirs(pathIN)
            
            year_limits = '{}to{}'.format(self.years[0], self.years[1])
            download_file = '{}-{}.nc'.format(self.variable, year_limits)
            download_path = pathIN.joinpath(download_file)
            download_form ={'product_type': self.product_type_ref,
                            'variable': self.variable,
                            'year': ['{}'.format(yy) for yy in range(self.years[0], self.years[1]+1)],
                            'month': ['{:02d}'.format(mm) for mm in range(self.start_months[0], self.start_months[1]+1)],
                            'time': '00:00',
                            'format': self.file_format_ref
                           }

            if 'pressure' in self.dataset_ref:
                download_form[levels_name] = self.levels

            #Download request
            self.down(download_form,
                      download_path,
                      dataset = self.dataset_ref
                     )

            
        toc = time.time()
        
        print('\nDownload elapsed time: {}.'.format(toc-tic))

            
