#!/usr/bin/env python
#
# Seasonal forecasts and hindcasts 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

import cds_download

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 check_hindcast_forecast(year):
    # Fucntion that determine if a request is for hindcast or forecast
    # It also returns the year as a string

    if year < 2017: 
        
        type_seasonal = 'seasonal_hindcast'
    
    else:
        
        type_seasonal = 'seasonal_forecast'
    
    return type_seasonal

class cds_request:

    def __init__(self,freq='',lvl='',month='',folder=''):
        # Frequency of the variable [daily, monthly]
        self.freq = freq
        # Month we want to download
        self.month = month
        # Single or Pressure Levels
        self.lvl = lvl
        # 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/seasonal'
        # 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 = ''

        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',
                'specific_humidity',
                'temperature',
                'u_component_of_wind',
                'v_component_of_wind'
            ]

        if c in ['slvl','SLVL','single']:
            # Single level variables
            return [
                '2m_dewpoint_temperature',
                '2m_temperature',
                'maximum_2m_temperature_in_the_last_24_hours',
                'mean_sea_level_pressure',
                'minimum_2m_temperature_in_the_last_24_hours',
                'sea_ice_cover',
                'sea_surface_temperature'
            ]

    def cds_form(self,var,year):
        # Function to define the CDS form depending on
        # the inline inputs defined when calling the script
        cds_form = {
            'originating_centre':'ecmwf',
            'system':'5',
            'variable':var,
            'year':str(year),
            'month':monthToNum(self.month),
            'format':'grib'
        }

        if self.lvl in ['plvl','PLVL','pressure']:
            
            cds_form['pressure_level'] = [
                '10','30','50',
                '100','200','300',
                '400','500','700',
                '850','925'
            ]
            
            self.lvl = 'PLVL'

        else:
            
            self.lvl = 'SLVL'

        if self.freq in ['monthly_mean','monthlymean']:

            cds_form['product_type'] = 'monthly_mean'
            
            self.freq = 'monthly_mean'
            
            cds_form['leadtime_month'] = [
                '1','2','3',
                '4','5','6'
            ]

        if self.freq == 'daily':

            cds_form['leadtime_hour'] = [
                '744',
                '1488',
                '2232'
            ]

            cds_form['day'] = '01'

        cds_type = ''

        if self.freq == 'daily' and self.lvl == 'SLVL':
            
            cds_type = 'seasonal-original-single-levels'

        if self.freq == 'daily' and self.lvl == 'PLVL':
            
            cds_type = 'seasonal-original-pressure-levels'

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

        if self.freq == 'monthly_mean' and self.lvl == 'PLVL':
            
            cds_type = 'seasonal-monthly-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:
                cds_download.cds_download(
                    self.cds_type,
                    self.form,
                    self.file_path,
                    url=self.cdsapiurl,
                    key=self.cdsapikey,
                    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)
            # ------------------------------------------------------

            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(1993,2020):

                    self.year = year

                    type_seasonal = check_hindcast_forecast(year)

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

                    g_name_list = [
                        type_seasonal,
                        self.form['originating_centre'],
                        'system'+self.form['system'],
                        self.freq,
                        self.lvl,
                        var,
                        str(year),
                        self.month
                    ]

                    if self.freq == 'daily':
                        grib_name = '_'.join(g_name_list)
                        grib_name += '_' + self.form['day'] + '.grib'
                    else:
                        grib_name = '_'.join(g_name_list) + '.grib'

                    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:]

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

        sys.exit(req_down.main())
    
    else:
        sys.exit()