#!/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, sys, gc, traceback, json

from dqc_wrapper.utils.utils_db import *
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_pdf import *

from dqc_wrapper.test00_file_format import FileFormatCheck
from dqc_wrapper.test01_standard_compliance import StandardComplianceTest
from dqc_wrapper.test02_spatial_completeness import SpatialCompletenessTest
from dqc_wrapper.test03_temporal_completeness import TemporalCompletenessTest
from dqc_wrapper.test04_spatial_consistency import SpatialConsistencyTest
from dqc_wrapper.test05_temporal_consistency import TemporalConsistencyTest
from dqc_wrapper.test06_valid_ranges import ValidRangesTest

"""
ECFLOW workflow source

[ ] Task 0 - Init/Create database, insert c3s datasets
[ ] Task 1 - Sync cds-c3s-forms
[ ] Task 2 - Download GRIB files + update database
[*] Task 3 - DQC checks + update database
[ ] Task 4 - PDF generation
[ ] Task 5 - Upload PDFs to drupal
[ ] Task 6 - Generate overview HTMLs
[ ] Task 7 - Delete and free up space
"""
def main(args):

    # Check parameters
    if len(args) != 1:
        print('Usage: python dqc_workflow.py <config_file>')
        return -1
    else:
        try:
            # General configuration
            with open(args[0]) as json_file:
                conf = json.load(json_file)
                dqc_conf_dir = conf['dqc_conf_dir']
                dqc_db_file = conf['dqc_db_file']
        except:
            print('ERROR: while reading the config file provided')
            exit(-1)

    # Create connection to SQLITE file
    try:
        conn = create_connection(dqc_db_file)
    except:
        print('ERROR: Could not connect to database')
        exit(-1)

    # LOOP 1: All datasets available in the database
    for d in select_datasets(conn):

        # Read configuration and check mandatory fields in the general section
        conf_path = os.path.join(dqc_conf_dir, d[0] + '.ini')
        c = CheckerConfig(conf_path)
        try:
            c.read_config()
        except Exception as e:
            print('ERROR: {}'.format(e))
            sys.exit()

        # LOOP 2: Every (file, data check) pair selected
        # Check only files to be checked (newly downloaded ones that are not checked)
        try:
            for ff,s in select_filedatachecks_todo(conn, d[0]):

                # Skip missing files
                file2check = os.path.join(c.path_files, ff+'.grib')
                ncfiles2check = os.path.join(c.path_files, ff + '.nc') # only a sample
                if not os.path.exists(file2check):
                    # Skip missing files
                    continue

                # Update file list to the current file
                c.file_list['original'] = [file2check ]
                c.file_list['nc_converter'] = [ ncfiles2check ]

                # Section header
                print_header("Data Quality Checker + ECFLOW [check={}]".format(s))

                # Free memory (every test)
                gc.collect()

                # Each section is a different test available
                test_res = None

                ## TEST 00 - File format (both for original GRIB + converted/experimental NetCDF)
                if s == 'file_format':
                    file_format_ck = FileFormatCheck(c.file_list['original'] + c.file_list['nc_converter'])
                    test_res = file_format_ck.main()
                    del file_format_ck

                ## TEST 010 - Standard compliance NetCDF (only one sample per dataset)
                elif s == 'standard_compliance_netcdf':
                    # TODO: define memory
                    standard_ck = StandardComplianceTest(c.file_list['nc_converter'],
                                                         c.forms_dir, c.cds_dataset, c.cds_variable, suffix='netcdf')
                    test_res = standard_ck.main()
                    del standard_ck

                ## TEST 011 - Standard compliance (original GRIB files)
                elif s == 'standard_compliance_grib':
                    standard_ck = StandardComplianceTest(c.file_list['original'],
                                                         c.forms_dir, c.cds_dataset, c.cds_variable, suffix='grib')
                    test_res = standard_ck.main()
                    del standard_ck

                ## TEST 02
                elif s == 'spatial_completeness':
                    spatial_completeness_ck = SpatialCompletenessTest(
                        c.file_list['original'], c.varname, c.datatype, c.cds_dataset, c.system, c.origin, c.masks_dir, c.land_sea_mask, c.land_variable
                    )
                    test_res = spatial_completeness_ck.main()
                    del spatial_completeness_ck

                ## TEST 03
                elif s == 'temporal_completeness':
                    temporal_completeness_ck = TemporalCompletenessTest(
                        c.res_dir, c.config_fname, c.file_list['original'], c.forms_dir,
                        c.cds_dataset, c.cds_variable, c.origin, c.system, c.is_seasonal, c.is_complete,
                        is_workflow=True
                    )
                    test_res = temporal_completeness_ck.main()
                    del temporal_completeness_ck

                ## TEST 04
                elif s == 'spatial_consistency':
                    spatial_consistency_ck = SpatialConsistencyTest(
                        c.file_list['original'], c.grid_type, c.grid_int, c.grid_map, c.dual_res
                    )
                    test_res = spatial_consistency_ck.main()
                    del spatial_consistency_ck

                ## TEST 05
                elif s == 'temporal_consistency':
                    temporal_consistency_ck = TemporalConsistencyTest(
                        c.res_dir, c.config_fname, c.file_list['original'], c.time_step, c.time_granularity, c.is_complete
                    )
                    test_res = temporal_consistency_ck.main()
                    del temporal_consistency_ck

                ## TEST 06
                elif s == 'valid_ranges':
                    valid_ranges_ck = ValidRangesTest(
                        c.res_dir, c.config_fname, c.file_list['original'], c.varname, c.valid_min, c.valid_max, c.datatype,
                        c.forms_dir, c.cds_dataset, c.cds_variable, c.cds_units
                    )
                    test_res = valid_ranges_ck.main()
                    del valid_ranges_ck

                else:
                    print('Skipping {}'.format(s))

                # Store test result (file + datacheck) in the database
                if test_res:
                    update_filedatacheck(conn, d[0], ff, s, test_res['res'], test_res['msg'])

        # Exceptions/Errors
        except KeyboardInterrupt:
            print_progress(100, 100)
            print('Cancelled, exiting...')
            exit()

if __name__ == "__main__":
    args = sys.argv[1:]
    args = [ '/data/joan/c3s512-wp1-datachecker/dqc_wrapper/conf/checker_workflow_conf.json' ]
    sys.exit(main(args))