#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)

import os
import re
import glob
import logging
import tempfile
import shutil
import zipfile
import json
from datetime import datetime, timedelta
import random
import string
import configparser

from decimal import Decimal
from contextlib import contextmanager
from datetime import datetime
import numpy as np

"""
Data checker configuration object
"""
class CheckerConfig():

    """
    Create an empty configuration object
    """
    def __init__(self, config_file):
        # Object
        self.config = configparser.RawConfigParser()
        self.config_path = config_file
        self.config_fname = os.path.basename(config_file)

        # General Attributes
        self.path_files = ""
        self.fpattern = ""
        self.log_dir = ""
        self.forms_dir = ""
        self.res_dir = ""
        self.masks_dir = ""
        self.log_file = ""

        # Dataset attributes
        self.cds_dataset = ""
        self.cds_variable = ""
        self.cds_units = ""
        self.varname = ""
        self.datatype = []

        # Test specific attributes
        self.mask_file = ""
        self.mask_var = ""
        self.origin = ""
        self.system = ""
        self.is_seasonal = ""
        self.grid_int = ""
        self.grid_type = ""
        self.grid_map = ""
        self.dual_res = ""
        self.time_step = ""
        self.time_granularity = ""
        self.valid_min = ""
        self.valid_max = ""

    """
    Read a configuration file into an object
    """
    def read_config(self, execute=True):

        # Read config from INI file
        if os.path.exists(self.config_path):
            self.config.read(self.config_path)
        else:
            raise Exception('Configuration file does not exist')

        # General settings
        self.path_files = self.check_param_config('general', 'input')
        self.fpattern = self.check_param_config('general', 'fpattern')
        self.log_dir = self.check_param_config('general', 'log_dir')
        self.res_dir = self.check_param_config('general', 'res_dir')
        self.forms_dir = self.check_param_config('general', 'forms_dir')
        self.masks_dir = self.check_param_config('general', 'masks_dir')
        self.cms_name = self.check_opt_param_config('general', 'cms_name', default='')

        # Dataset settings
        self.cds_dataset = self.check_opt_param_config('dataset', 'cds_dataset')
        self.cds_variable = self.check_opt_param_config('dataset', 'cds_variable')
        self.cds_units = self.check_opt_param_config('dataset', 'cds_units')
        self.varname = self.check_opt_param_config('dataset', 'variable')
        self.datatype = self.check_datatype_config('dataset', 'datatype')
        self.is_complete = self.check_opt_param_config('dataset', 'is_complete', default='yes')
        self.origin = self.check_opt_param_config('dataset', 'origin', default='')
        self.system = self.check_opt_param_config('dataset', 'system', default='')

        # Tests configuration
        for s in self.config.sections():

            # Each section has different parameters
            try:
                ## TEST 02
                if s == 'spatial_completeness':
                    self.land_sea_mask = self.check_opt_param_config(s, 'land_sea_mask', default='no')  # default is without mask
                    self.land_variable = self.check_opt_param_config(s, 'land_variable', default='no')  # default is ocean_variable
                ## TEST 03
                elif s == 'temporal_completeness':
                    self.is_seasonal = self.check_opt_param_config(s, 'is_seasonal', default='no')  # default is not seasonal
                ## TEST 04
                elif s == 'spatial_consistency':
                    self.grid_int = self.check_param_config(s, 'grid_interval')
                    self.grid_type = self.check_param_config(s, 'grid_type')
                    self.grid_map = self.check_opt_param_config(s, 'grid_mapping_name', default='')  # default is empty
                    self.dual_res = self.check_opt_param_config(s, 'dual_resolution', default='no')  # default is one resolution
                ## TEST 05
                elif s == 'temporal_consistency':
                    self.time_step = self.check_param_config(s, 'time_step')
                    self.time_granularity = self.check_param_config(s, 'time_granularity')
                ## TEST 06
                elif s == 'valid_ranges':
                    self.valid_min = self.check_param_config(s, 'valid_min')
                    self.valid_max = self.check_param_config(s, 'valid_max')
            except:
                raise Exception('Reading parameters from configuration file')

        # Run mode vs dry run
        if execute:
            # Directory checks
            if not os.path.exists(self.path_files):
                raise Exception('Input directory does not exist. Check configuration file section [general]')
            if not self.log_dir:
                raise Exception('Log directory not defined. Check configuration file section [general], value log_dir')
            if not self.res_dir:
                raise Exception('Results directory not defined. Check configuration file section [general], value res_dir')
            if not os.path.exists(self.masks_dir):
                raise Exception('Land/Sea Masks directory not defined or does not exist. Check configuration file section [general], value res_dir')
            if not os.path.exists(self.res_dir):
                os.makedirs(self.res_dir)

            # Init logging and list files
            self.log_file = initiate_logging(self.log_dir, self.path_files, self.config_fname)
            self.file_list = dir_or_file(self.path_files, self.fpattern)

            # Not needed anymore
            self.config.remove_section('general')
            self.config.remove_section('dataset')

    """
    Check a MANDATORY parameter inside the config file (based on RawConfigParser)
    """
    def check_param_config(self, section, parameter):
        try:
            return self.config[section][parameter].strip()
        except:
            raise Exception('Parameter {p} not defined. Check configuration file section [{s}]'.format(s=section, p=parameter))

    """
    Check an OPTIONAL parameter inside the config file (based on RawConfigParser)
    """
    def check_opt_param_config(self, section, parameter, default=''):
        try:
            return self.config[section][parameter].strip()
        except:
            return default

    """
    Check datatype parameter in the config. Can be an array of CSV.
    """
    def check_datatype_config(self, section, parameter, default=[]):
        try:
            dt = [ d.strip() for d in self.config[section][parameter].split(',') ]
            if len(dt) == 1 and dt[0] == '':
                dt = []
            return dt
        except:
            return default

    """
    Save a config section/test after a series of tests, log output to console
    """
    def save_config(self, section, test_res, verbose=True):

        if test_res is not None:
            if not self.config.has_section(section + '_result'):
                self.config.add_section(section + '_result')
            self.config.set(section + '_result', 'res', test_res['res'])
            if test_res['msg'] is not None:
                self.config.set(section + '_result', 'msg', test_res['msg'].replace('%', 'percent'))

            if verbose:
                print()
                if test_res['res'] == 'ok':
                    print('INFO: {}'.format(test_res['msg']))
                elif test_res['msg'] is None:
                    print('INFO: Not available yet')
                elif test_res['res'] == 'warn':
                    print('WARN: {}'.format(test_res['msg']))
                else:
                    print('ERROR: {}'.format(test_res['msg']))

"""
Load text file
"""
def load_text_file(path):
    try:
        with open(path, 'r') as fp:
            return fp.read()
    except:
        return "ERROR: Filename {} does not exists".format(path)

"""
Print a header before a progress start
"""
def print_header(text):
    string = "\n" + text + "\n"
    string += "." + "." * 80 + "\n"
    print(string)

"""
Returns a temporary dir that cleans up on exit (for unzipping)
"""
@contextmanager
def tempdir():
    path = tempfile.mkdtemp()    
    try:
        yield path
    finally:
        try:
            shutil.rmtree(path)
        except IOError:
            raise Exception('Failed to clean up temp dir {}'.format(path))   


"""
Return file list matching a given pattern
"""
def dir_or_file(abspath, pattern, quiet=False):

    # Single file
    if os.path.isfile(abspath):
        print("INFO: Checking file")
        print(abspath)
        list_files = list_files_nc = [ abspath ]

    # Directory
    elif os.path.isdir(abspath):        
        # Empty pattern equals all files
        if pattern.rstrip() == "":
            pattern = '*'

        # List all files in the folder (grib, netcdf-converter)
        list_files = glob.glob(os.path.join(abspath, pattern))
        list_files_nc = glob.glob(os.path.join(abspath, pattern.replace('.grib', '.nc')))

        # Sort by filename / date
        list_files = sorted(list_files)

        # Output
        if not(quiet):
            print("INFO: Path {}".format(abspath))
            print("INFO: Found {} files to check".format(len(list_files)))
            if len(list_files):
                print("INFO: [1]: {}".format(os.path.basename(list_files[0])))
                print("INFO: [N]: {}\n".format(os.path.basename(list_files[-1])))
            else:
                raise Exception('No files found for analysis')

        # Logging files to analyze
        basenames = []
        for f in list_files:
            basenames.append(os.path.basename(f))

    return {
        'original': list_files,
        'nc_converter': list_files_nc
    }

"""
Format float function, number of decimals for floats is set to 3 by default
"""
def format_float_single(ff, decimals=3):
    try:
        ff = float(ff)
        if abs(ff) > 0.1:
            nr = "{0:." + str(decimals) + "f}"
            return nr.format(ff)
        else:
            return '%.2E' % Decimal(ff)
    except:
        return 'N/A'

"""
Format float function
"""
def format_float(ff, decimals=3):

    if type(ff) == list or type(ff) == np.ndarray:
        str='['
        for e in ff:
            str = str + format_float_single(e, decimals) + ', '

        return str[:-1] + ']'
    else:
        return format_float_single(ff, decimals)

"""
Get zip name
"""
def get_zip_fname(res_dir, config_name):

    # Output zip file name
    zipfname = os.path.join(res_dir, os.path.basename(config_name) + '.zip')
    if '.ini' in config_name:
        zipfname = os.path.join(res_dir, os.path.basename(config_name).replace('.ini','.zip'))
    
    return zipfname

"""
Zip a list of files 
"""
def zip_results(zipfname, file_list):

    # All checks must be passed, all files must exist
    for f in file_list:
        if not (os.path.exists(f)):
            return

    # Zip a list of files
    with zipfile.ZipFile(zipfname, 'w') as zipf:        
        for f in file_list:
            zipf.write(f, os.path.basename(f), compress_type=zipfile.ZIP_DEFLATED)
    print('INFO: Saving zip to {}'.format(zipfname))

"""
Schedule the delete of already processed files by the checker
a crontab job will check the file and process the content
todo: add config section to enable/disable this
"""
def schedule_datafiles_deletion(out_dir, file_list):

    try:
        #load settings
        this_dir = os.path.dirname(__file__)
        script_settings = os.path.join(this_dir, "../conf/scripts_settings_conf.json")
        with open(script_settings) as json_file:
            conf = json.load(json_file)
            output_folder = conf['dqc_files_removal_path']
            enabled_file_deletion = conf['dqc_files_removal_enable']
            #load defaults in case not defined
            if not output_folder :
                output_folder = out_dir
            if not enabled_file_deletion :
                enabled_file_deletion = False

        if enabled_file_deletion :
            #scheduling will take place a week after the data/QARs has been processed
            current_date = datetime.now().strftime("%Y/%m/%d")
            date = datetime.strptime(current_date, "%Y/%m/%d")
            deletion_date = date + timedelta(days=7)
            # name convention datafiles-to-be-deleted-dd-mm-yyyy.checked
            out_file = os.path.join(output_folder, "datafiles-to-be-deleted-" + datetime.strftime(deletion_date, "%Y%m%d") + ".checked")
            # Update if exists
            print('INFO: generating scheduled checked files archive {}'.format(out_file))
            to_be_deleted_file = open(out_file, 'w')
            if os.path.exists(out_file):
                for f in file_list:
                    # append the files to the archive at the end
                    print(f)
                    to_be_deleted_file.write(f)
            to_be_deleted_file.close()
    except Exception as e:
        print('ERROR: generating scheduled checked files archive')
        print(str(e))


"""
Logging reset
"""
def logging_reset():
    # Remove all handlers associated with the root logger object.
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

"""
Create logging associated to a test [00-0X]
"""       
def initiate_logging(path_logs, path_files, testid):

    # Directory check
    if not os.path.exists(path_logs):
        os.makedirs(path_logs)

    # Log file name
    default_log_name = '_'.join(
        [
            'LOG',
            testid,
            datetime.now().strftime('%Y%m%d_%H%M%S')
        ]
    )
    default_path = os.path.join(path_logs, default_log_name + '.txt')
    print('INFO: Saving log in {}'.format(default_path))

    # Logging configuration
    logging.basicConfig(
        filename=default_path,
        format='%(name)s:%(levelname)s:%(asctime)s:%(message)s', 
        level=logging.INFO,
        filemode = 'w'
    )

    # Write header
    stringA =  '\n' + '#'*80 + '\n'
    stringA += '\n RUNNING TEST {}'.format(testid)
    stringA += '\n DIR = ' + path_files + '\n'
    stringA += '\n' + '#'*80
    logging.info(stringA)

    return default_path

"""
Write a result of a check to a file
"""
def save_result(config, res_dir, fname):

    # Save result
    out_file = os.path.join(res_dir, fname)

    # Update if exists
    if os.path.exists(out_file):
        try:
            old_config = configparser.RawConfigParser()
            old_config.read(out_file)
            for s in old_config.sections():
                if config.has_section(s):
                    for k,v in config[s].items():
                        old_config.set(s, k, v)

            # Only if content detected
            if old_config.sections() != []:
                config = old_config
                print('INFO: Updating {}'.format(out_file))
        except:
            pass # unreadable config

    # Write object to disk
    with open(out_file, 'w') as configfile:
        # Result
        print('INFO: Saving result to {}'.format(out_file))
        config.write(configfile)

    return out_file

"""
Progress bar function
"""
def print_progress(iteration, total, prefix='', suffix='', decimals=1, length=80, fill='#', printEnd="\r"):

    # Call in a loop
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filledLength = int(length * iteration // total)
    bar = fill * filledLength + '-' * (length - filledLength)
    print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end = printEnd)
    
    # Print New Line on Complete
    if iteration == total: 
        print('\n')

"""
Get dates from file name list
"""
def get_dates_from_files(file_list):

    dates = []
    for f in file_list:
        pt = re.search("([0-9]{8})", f) # YYYYMMDD
        try:
            dates.append(datetime.strptime(pt[0], '%Y%m%d'))
        except:
            raise Exception('Invalid date format. Please name files accordingly: YYYYMMDD')
    return dates

"""
Filter out dates according to file names
"""
def filter_dates(dates_list, file_list):

    ret_dates = []
    for d in dates_list:
        for dd in get_dates_from_files(file_list):
            if dd.year == d.year and dd.month == d.month:
                ret_dates.append(d)

    return ret_dates

"""
Get a pattern from a file (filename without the date)
"""
def file_pattern(file):
    pt = re.split("([0-9]{8})", file)  # YYYYMMDD
    return pt[0]

"""
Cleanup directory and files
"""
def cleanup_directory(dir_path):
    try:
        if not(dir_path is None):
            shutil.rmtree(dir_path)
    except OSError as e:
        raise Exception("Error: %s : %s" % (dir_path, e.strerror))

"""
Cleanup files
"""
def cleanup_files(files):
    try:
        for f in files:
            os.remove(f)
    except OSError as e:
        raise Exception("Error: %s : %s" % (f, e.strerror))

"""
Combine two test results (append messages and decide ok)
"""
def combine_tests(test_res1, test_res2):

    # OK/WARN/ERR: only if two are ok
    combined = test_res1['res'].strip() + ',' + test_res2['res'].strip()
    if 'err' in combined:
        res = 'err'
    elif 'warn' in combined:
        res = 'warn'
    else:
        res = 'ok'

    # Simply append messages to a paragraph
    msg = test_res1['msg'] + '. ' + test_res2['msg']

    # Return JSON object
    return {
            'res': res,
            'msg': msg
    }

"""
Generate check result and message (used in the workflow / database)
"""
def overall_result(checkId, checkResult, currentResult, checkMsg, currentMsg):

    ## ERR - it does not matter, error from now on
    if checkResult == -1:
        res = 'err'
    ## WARN - only an error prevails
    elif checkResult == 2:
        res = 'warn'
        if currentResult == -1:
            res = 'err'
    ## OK - we keep the result unless it is ok
    elif checkResult == 1:
        if currentResult == 1:
            res = 'ok'
        else:
            res = currentResult
    ## UNDEF
    else:
        res = 0

    ## MSG
    if checkMsg == "":
        msg = currentMsg
    else:
        msg = checkMsg

    return {
            'res': res,
            'msg': msg
    }


"""
Rename a file according to naming convention
example: 'reanalysis-era5-single-levels-monthly-means-volumetric_soil_water_layer_2.ini'
result: reanalysis_era5_single_levels_monthly_means-era5-volumetric_soil_water_layer_2-valid_ranges.html
"""
def rename_according_convention(filename, testname, extension):

    filename = filename.replace('.ini', '')

    if 'seasonal' in filename:
        centre = filename.split('-')[-2]
        system = filename.split('-')[-1]
        var = filename.split('-')[-3]
        dataset = ('_').join(filename.split('-')[:-3])
        filename = '{dataset}-{originating_centre}-{system}-{variable}-{check}.{ext}'.format(dataset=dataset,
                                                                                    originating_centre=centre,
                                                                                    system=system, variable=var,
                                                                                    check=testname, ext=extension)
    if 'reanalysis' in filename:
        var = filename.split('-')[-1]
        
        if 'uerra' in filename:
            system = filename.split('-')[-2]
            dataset = ('_').join(filename.split('-')[:-2])
            filename = '{dataset}-{system}-{variable}-{check}.{ext}'.format(dataset=dataset,
                                                                            system=system, variable=var,
                                                                            check=testname, ext=extension)
        
        if 'era5' in filename:
            dataset = ('_').join(filename.split('-')[:-1])

            if 'era5-land' in filename:
                filename = '{dataset}-era5_land-{variable}-{check}.{ext}'.format(dataset=dataset,
                                                                                variable=var,
                                                                                check=testname, ext=extension)

            else: # era5
                filename = '{dataset}-era5-{variable}-{check}.{ext}'.format(dataset=dataset,
                                                                            variable=var,
                                                                            check=testname, ext=extension)

    return filename

"""
Read file where every line is : variable_name, shortName
and return list with all shortNames
"""
def get_varname_list(path):
    _list = []
    for line in load_text_file(path).splitlines():
        _list.append(line.split(',')[1])
    return _list



"""
Generate log file from dqc_workflow stored messages
"""
def generate_workflow_log():
    log="LOG_seasonal-monthly-single-levels-2m_temperature-ecmwf-5.ini_20200724_152025.txt"