#!/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
from decimal import Decimal
from contextlib import contextmanager
from datetime import datetime


"""
Check a MANDATORY parameter inside the config file (based on ConfigParser)
"""
def check_param_config(config, section, parameter):
    try:
        return 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(config, section, parameter, default=''):
    try:
        return config[section][parameter].strip()       
    except:                
        return default

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

"""
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):
    if os.path.isfile(abspath):
        print("INFO: Checking file")
        print(abspath)
        # do the file checking
        return [abspath]

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

        # list all files in the folder
        list_files = glob.glob(os.path.join(abspath, pattern))
        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 list_files            

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

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

    print()
    if test_res != None:
        config.add_section(section + '_result')
        config.set(section + '_result', 'res', test_res['res'])
        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']))

    return config

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