#!/usr/bin/env python
#
# Reanalysis ERA5 downloader
#
# Barcerlona Supercomputing Center - Centro Nacional de Supercomputacion
#        Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: Ivan Jose Cernicharo Ortiz (ivan.cernicharo@bsc.es)
#              Jesus Pena-Izquierdo (jesus.pena@bsc.es)
#

import os
import sys
import logging
import time
import json

from C3S512.cds import downloader

def monthToNum(Month):
    # Function to convert a month as a string to a string of two digits    
    months_dict = {
        'january' : '01',
        'february' : '02',
        'march' : '03',
        'april' : '04',
        'may' : '05',
        'june' : '06',
        'july' : '07',
        'august' : '08',
        'september' : '09', 
        'october' : '10',
        'november' : '11',
        'december' : '11'
    }

    return months_dict[Month]

def return_day():
    list_days = [
        '01','02','03',
        '04','05','06',
        '07','08','09',
        '10','11','12',
        '13','14','15',
        '16','17','18',
        '19','20','21',
        '22','23','24',
        '25','26','27',
        '28','29','30',
        '31'
    ]
    return list_days


class cds_request:

    def __init__(self,freq = None,lvl = None,month = None,folder = None,format = None,prod_type = None):

        # Frequency of the variable [daily, monthly]
        self.freq = freq
        # Month we want to download
        self.month = month
        # Single or Pressure Levels
        self.lvl = lvl
        # Product type
        self.prod_type = prod_type
        # Format file
        self.format = format.lower()
        # Folder where the files are supossed to be stored
        self.fold = folder
        # Folder where the file is supossed to be stored. WIll be modified for each variable
        self.fold1 = folder
        # Download form
        self.form = ''
        # File path
        self.file_path = ''
        # Log path
        self.log_path = ''
        # Default LOG path
        self.fold2 = 'Python_LOGs/era5'
        # Default credentials json for cdsapi
        # This file should be stored in the same folder where this script is
        self.creds_file = 'cdsapi_icernich.json'
        # Database name
        self.cds_type = ''
        # System, by default ERA5
        self.system = 'era5'


        # if not os.path.exists(self.creds_file):
        #     # checking if the credentials file exists            
        #     string = '\n\tCDS API credentials file not found\n'
            
        #     logging.error(string)
            
        #     sys.exit()

        # self.cdsapiurl,self.cdsapikey = self.get_cdsapi_credentials()

    def get_cdsapi_credentials(self):
        # Function to read the credentials JSON
        with open(self.creds_file) as j_doc:
            
            cds_creds = json.load(j_doc) 
        
        j_doc.close()

        return cds_creds['url'],cds_creds['key']

    def list_vars(self):
        #Function to return a list of variables to download        
        c = self.lvl
        
        # if c in ['plvl','PLVL','pressure']:
        #     # Pressure levels variables
        #     return [
        #         'geopotential',
        #         'temperature',
        #         'u_component_of_wind',
        #         'v_component_of_wind',
        #         'vertical_velocity',
        #         'relative_humidity'
        #     ]

        if c in ['slvl','SLVL','single']:
            # Single level variables
            # return [
            #     '10m_u_component_of_wind',
            #     '10m_v_component_of_wind',
            #     '2m_dewpoint_temperature',
            #     '2m_temperature',
            #     'mean_sea_level_pressure',
            #     'total_precipitation'
            # ]
            return [
                '10m_u_component_of_wind'
            ]
        
    def cds_form(self,var,year):
        # Function to define the CDS form depending on
        # the inline inputs defined when calling the script
        cds_form = {}

        if self.freq == 'monthly_mean' and self.prod_type == None:
            cds_form['product_type'] = 'monthly_averaged_reanalysis'                

        if self.freq == 'daily':
            # By default
            cds_form['day'] = '01'

            if  self.prod_type == None:
                cds_form['product_type'] = 'reanalysis'

        cds_form['variable'] = var
        cds_form['month'] = monthToNum(self.month)
        cds_form['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'
        ]
        cds_form['year'] = str(year)
        cds_form['format'] = self.format

        if self.lvl in ['plvl','PLVL','pressure']:
            
            cds_form['pressure_level'] = [
                '1','2','3','5','7','10',
                '20','30','50','70','100','125',
                '150','175','200','225','250','300',
                '350','400','450','500','550','600',
                '650','700','750','775','800','825',
                '850','875','900','925','950','975',
                '1000'
            ]
            
            self.lvl = 'PLVL'

        else:
            
            self.lvl = 'SLVL'

        if self.freq == 'monthly_mean' and self.lvl == 'SLVL':
            cds_type = 'reanalysis-era5-single-levels-monthly-means'

        if self.freq == 'monthly_mean' and self.lvl == 'PLVL':
            cds_type = 'reanalysis-era5-pressure-levels-monthly-means'

        if self.freq == 'daily' and self.lvl == 'SLVL':
            cds_type = 'reanalysis-era5-single-levels'

        if self.freq == 'daily' and self.lvl == 'PLVL':
            cds_type = 'reanalysis-era5-pressure-levels'
            
        return cds_form,cds_type

    def cds_download(self):
        # Function that make the request to the CDS to retrieve data
        if not os.path.exists(self.file_path):
            
            # Information for the LOG file to store the CDS form
            separators = '\n\t' + "-"*10 + " CDSAPI FORM " + "-"*10
            
            folderf = '/'.join(self.file_path.split('/')[:-1])
            filef = self.file_path.split('/')[-1]

            string = separators + '\n'
            string += "Data set name: " + self.cds_type + '\n'
            string += "Path to file: " + folderf + '\n'
            string += "File name: " + filef + '\n'
            string += json.dumps(self.form,indent=4,sort_keys=True)
            string += '\n' + separators
            
            logging.info(string)
            # ------------------------------------------------------
            
            dtime0 = time.time()

            # try:
            #     downloader(
            #         self.cds_type,
            #         self.form,
            #         self.file_path,
            #         url=self.cdsapiurl,
            #         key=self.cdsapikey,
            #         user='icernich'
            #     )

            try:
                downloader(
                    self.cds_type,
                    self.form,
                    self.file_path,
                    user='icernich'
                )

            except Exception as e:

                # Information for the LOG file when an ERROR ocurrs
                separators = "\n\t\t" + "-"*10 + " ERROR " + "-"*10 + "\n"

                logging.error(separators)
                logging.error(e)
                logging.error(separators)

                logging.error(string)
                # ------------------------------------------------------

            # Information for the LOG file for TIME elapsed since the request
            # was made until the data has been downloaded
            dtime1 = time.time()
            
            difftime = dtime1 - dtime0

            str_time = str(int(difftime//3600)) + ' hours, '
            str_time += str(int((difftime%3600)//60)) + ' minutes, '
            str_time += str(round((difftime%3600)%60,3)) + ' seconds '

            separator = '\n\t'+ "-"*10 + " TIME ELAPSED " + "-"*10 

            string = separator + '\n'
            string += '\n\t\t%s\n' %str_time
            string += separator

            logging.info(string)
            # ------------------------------------------------------
        
        else:

            folderf = '/'.join(self.file_path.split('/')[:-1])
            filef = self.file_path.split('/')[-1]

            # Information for the LOG file if the file already exists
            string = '\n' + '##########################'*3 + '\n'
            string += 'Requested file\n\t%s' %folderf
            string += '\nIs already in directory\n\t%s' %filef
            string += '\nCDS API FORM:\n'
            string += json.dumps(self.form,indent=4,sort_keys=True)
            string += '\n' + '##########################'*3 + '\n'

            logging.warning(string)
            # ------------------------------------------------------

    def main(self):
        # This function controls the main routine to download data from the CDS

        conditional1 = self.freq != '' and self.month != ''

        conditional2 = self.lvl != '' and self.fold1 != ''

        if conditional1 and conditional2:

            if not os.path.exists(self.fold2):

                os.makedirs(self.fold2)
            
            LOGname = 'LOG_' + sys.argv[0].replace('.py','')

            LOGname += '_' + self.freq + '_' + self.lvl + '_' + self.month + '.txt'

            self.log_path = os.path.join(self.fold2, LOGname)

            logging.basicConfig(
                filename=self.log_path, 
                format='%(levelname)s:%(asctime)s:%(message)s', 
                level=logging.INFO
            )
            
            # Information for the LOG file about the dataset and the input parameters
            string = "\n" + "#"*90 + '\n' + "#"*90 + '\n#'
            string += '\t%s %s on %s levels for %s\n#' %(self.cds_type,self.freq,self.lvl,self.month)
            string += "\n" + "#"*90

            logging.info(string)
            # ------------------------------------------------------

            # Checking if file format is grib or netcdf
            if not self.format in ['grib','netcdf']:
                string = '\n' + '########## ERROR #########'*3 + '\n'
                string += 'Requested file format\n\t%s' %self.format
                string += '\nIs not correct\n'
                string += '\n' + '########## ERROR #########'*3 + '\n'

                logging.error(string)

                sys.exit()

            for var in self.list_vars():

                self.var = var

                self.fold1 = os.path.join(self.fold,self.freq,self.lvl,str(var))

                for year in range(1979,2019):

                    if self.freq == 'daily':

                        self.year = year

                        self.form,self.cds_type = self.cds_form(var,year)

                        self.form['day'] = return_day()

                        g_name_list = [
                            self.system,
                            self.lvl,
                            self.form['product_type'],
                            var,
                            str(year),
                            self.month
                        ]

                        grib_name = '_'.join(g_name_list)

                        if self.format == 'grib':
                                grib_name += '.grib'

                        if self.format == 'netcdf':
                            grib_name += '.nc'

                        self.file_path = os.path.join(self.fold1, grib_name)

                        # Information for the LOG file on the variable name and year
                        sep = '_*_*_*_*_*_*_*_*_*_*_*_*_*'*3
                        string = '\n' + sep + '\n'
                        string += '\tDownloading variable %s for %s \n' %(var,str(year))

                        logging.info(string)
                        # ------------------------------------------------------

                        self.cds_download()
                    

if __name__ == '__main__':

    args = sys.argv[1:]

    # Needs at least 5 arguments. & is the max number of input arguments

    if len(args) == 6:
        
        # Inputs for the download routine. Must have 6 inputs at least to make in work
        freq,lvl,month,folder,prod_type,format = args[0],args[1],args[2],args[3],args[4],args[5]


        req_down = cds_request(freq=freq,lvl=lvl,month=month,folder=folder,format=format,prod_type=prod_type)

        sys.exit(req_down.main())

    elif len(args) == 5:

        # Inputs for the download routine. Must have 5 inputs at least to make in work
        freq,lvl,month,folder,format = args[0],args[1],args[2],args[3],args[4]


        req_down = cds_request(freq=freq,lvl=lvl,month=month,folder=folder,format=format)

        sys.exit(req_down.main())

    
    else:
        sys.exit()