#!/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 random
import string
import configparser

from decimal import Decimal
from contextlib import contextmanager
from datetime import datetime

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

    """
    Create an empty configuration object
    """
    def __init__(self, config_file):
        # Object
        self.config = configparser.ConfigParser()
        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.tmp_dir = ""
        self.log_file = ""

        # Dataset attributes
        self.cds_dataset = ""
        self.cds_variable = ""
        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):

        # 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.tmp_dir = self.check_param_config('general', 'tmp_dir')

        # 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.varname = self.check_opt_param_config('dataset', 'variable')
        self.datatype = self.check_datatype_config('dataset', 'datatype')

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

            # Each section is a different test available
            try:
                ## TEST 02
                if s == 'spatial_completeness':
                    self.mask_file = self.check_opt_param_config(s, 'mask_file', default='')  # default is without mask
                    self.mask_var = self.check_opt_param_config(s, 'mask_var', default='lsm') # default is lsm (land_sea_mask)
                ## TEST 03
                elif s == 'temporal_completeness':
                    self.origin = self.check_param_config(s, 'origin')
                    self.system = self.check_param_config(s, 'system')
                    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')

        # Splitting dates if datasets are heavy is totally optional
        self.split_dates = self.check_opt_param_config('dataset', 'split_dates', default='no')

        # 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.res_dir):
            os.makedirs(self.res_dir)
        if not os.path.exists(self.tmp_dir):
            os.makedirs(self.tmp_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, self.tmp_dir, self.split_dates, self.config_fname)

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

    """
    Check a MANDATORY parameter inside the config file (based on ConfigParser)
    """
    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 ConfigParser)
    """
    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:
            return [ d.strip() for d in self.config[section][parameter].split(',') ]
        except:
            return default

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

        print()
        if test_res != None:
            self.config.add_section(section + '_result')
            self.config.set(section + '_result', 'res', test_res['res'])
            self.config.set(section + '_result', 'msg', test_res['msg'].replace('%', 'percent'))

            if test_res['res'] == 'ok':
                print('INFO: {}'.format(test_res['msg']))
            elif test_res['res'] == 'warn':
                print('WARN: {}'.format(test_res['msg']))
            else:
                print('ERROR: {}'.format(test_res['msg']))

"""
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))   

"""
Splits a list of files making them smaller (into a temporary directory)
- For seasonal - stepRange
- For reanalysis - isobaricInhPa
"""
def grib_split(list_files, tmp_dir):

    # Temporary rules file
    rules_file = os.path.join(tmp_dir, 'rules')
    script_file = os.path.join(tmp_dir, 'split.sh')
    with open(rules_file, 'w') as fp1, open(script_file, 'w') as fp2:
        fp1.write('write "{}/[isobaricInhPa]_[stepRange]_[dataDate].grib";\n'.format(tmp_dir))

        # Split all files and return dir
        for grib_file in list_files:
            fp2.write('/usr/bin/grib_filter {r} {g};\n'.format(r=rules_file, g=grib_file))

    print('Filtering ...')
    os.system('/bin/bash {} 2>&1'.format(script_file))

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

    # Single file
    tmp_dir_split = None
    if os.path.isfile(abspath):
        print("INFO: Checking file")
        print(abspath)
        list_files = list_files_split = 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, split(optional))
        list_files = glob.glob(os.path.join(abspath, pattern))
        list_files_nc = glob.glob(os.path.join(abspath, pattern.replace('.grib', '.nc')))
        if split_dates == 'no':
            list_files_split = list_files
        else:
            random_str = ''.join(random.choices(string.ascii_uppercase + string.digits, k=6))
            tmp_dir_split = os.path.join(tmp_dir, config_fname + '_' + random_str)
            if not(os.path.exists(tmp_dir_split)):
                os.makedirs(tmp_dir_split)
            grib_split(list_files, tmp_dir_split)
            list_files_split = glob.glob(os.path.join(tmp_dir_split, '*.grib'))

        # 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,
        'split': list_files_split,
        'tmp_dir_split': tmp_dir_split
    }

"""
Format float function
"""
def format_float(ff):
    if ff > 0.1:
        return "{0:.2f}".format(float(ff))
    else:
        return '%.2E' % Decimal(float(ff))

"""
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))

"""
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)
    print('INFO: Saving result to {}'.format(out_file))
    with open(out_file, 'w') as configfile:
        # Result
        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

"""
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))

"""
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
    }