#!/usr/bin/env python3

print('''
filter_json.py - a script from the AQUAMIS pipeline:
https://gitlab.com/bfr_bioinformatics/AQUAMIS.git
''')

# %% Packages #################################################################

# Python Standard Packages
import sys
import os
import logging
import traceback
import argparse
import json
from pprint import pprint

# Other Packages
import cerberus
import pandas

from helper_functions import *


# %% Exception Handling #######################################################

def handle_exception(exc_type, exc_value, exc_traceback):
    if issubclass(exc_type, KeyboardInterrupt):
        sys.__excepthook__(exc_type, exc_value, exc_traceback)
        return
    logging.error(''.join(["Uncaught exception: ", *traceback.format_exception(exc_type, exc_value, exc_traceback)]))


# Install exception handler
sys.excepthook = handle_exception


# %% Define Source of Arguments ###############################################

def parse_arguments():
    repo_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
    parser = argparse.ArgumentParser()

    # input
    parser.add_argument('--json_in', '-i', help = 'Path to the JSON file of combined BfR pipeline results, e.g. sample.merge.json',
                        type = os.path.abspath, required = True)
    # output
    parser.add_argument('--json_out', '-o', help = 'Path to the JSON file of filtered results in the BeONE format, e.g. sample.beone.json',
                        type = os.path.abspath, required = True)
    # params
    parser.add_argument('--sample', '-s', help = 'String of the Sample name',
                        required = True)
    parser.add_argument('--validation', '-v', help = 'Path to the BeONE JSON schema for Validation, default: ' + os.path.join(repo_dir, 'resources', 'BfR-GenSurv_schema_v20220619_full.json'),
                        type = os.path.abspath, required = True)
    parser.add_argument('--filter', '-f', help = 'Path to the BeONE JSON schema for Filtering, default: ' + os.path.join(repo_dir, 'resources', 'BeONE_schema_filter_v20220619.json'),
                        type = os.path.abspath, required = True)
    # log
    parser.add_argument('--log', '-l', help = 'Path for the logfile',
                        type = os.path.abspath, required = True)

    return parser.parse_args()


# Create Arguments for Debugging
def create_arguments():
    repo_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
    debug_sample = 'SRR498433'  # 'GMI-17-001-DNA' 'Se-Germany-BfR-0001' 'Se-BfR-0001' 'SRR498433'
    debug_workdir = os.path.join(repo_dir, "test_data")  # 'test_data' 'test_data_beone' 'test_data_workshop'

    args = argparse.Namespace(
        json_in = os.path.join(debug_workdir, 'json', 'post_qc', debug_sample + '.aquamis.json'),
        json_out = os.path.join(debug_workdir, 'json', 'filtered', debug_sample + '.aquamis.json'),
        sample = debug_sample,
        validation = os.path.join(repo_dir, 'resources', 'AQUAMIS_schema_v20230906.json'),
        filter = os.path.join(repo_dir, 'resources', 'AQUAMIS_schema_filter_v20230906.json'),
        log = os.path.join(debug_workdir, 'logs', debug_sample + '_filterJson_debug.log')
    )
    return args


# Translate ArgParse to Snakemake Object
def arguments_to_snakemake(args):
    global snakemake
    snakemake_input_df = pandas.DataFrame(
        data = [['json_in', args.json_in]],
        columns = ['index', 'input'])
    snakemake_output_df = pandas.DataFrame(
        data = [['json_out', args.json_out]],
        columns = ['index', 'output'])
    snakemake_params_df = pandas.DataFrame(
        data = [['sample', args.sample],
                ['validation', args.validation],
                ['filter', args.filter]],
        columns = ['index', 'params'])
    snakemake_log_df = pandas.DataFrame(
        data = [['log', args.log]],
        columns = ['index', 'log'])
    snakemake_df = snakemake_input_df.merge(snakemake_output_df.merge(snakemake_params_df.merge(snakemake_log_df, how = 'outer'), how = 'outer'), how = 'outer')
    snakemake = snakemake_df.set_index(snakemake_df['index'])


# Define Source of Arguments
def python_or_snakemake(debug: bool):
    try:
        snakemake.input[0]
    except NameError:
        if debug:
            arguments_to_snakemake(args = create_arguments())
        else:
            arguments_to_snakemake(args = parse_arguments())


# %% Import Data ##############################################################

def import_data():
    global json_sample
    global json_schema_validation
    global json_schema_filter

    # Import Sample
    with open(snakemake.input['json_in'], 'r') as file:
        json_sample = json.load(fp = file)  # without arg object_pairs_hook=unquote_hook to mimic ExtendsClasses schema

    # Import Schemas
    with open(snakemake.params['validation'], 'r') as file:
        json_schema_validation = json.load(fp = file)

    with open(snakemake.params['filter'], 'r') as file:
        json_schema_filter = json.load(fp = file)


# %% Validate and Filter Sample ###############################################

def validate_sample(json_sample, json_schema_validation):
    # Cerberus: Create Validator - Convert custom JSON schema to Cerberus schema
    cerberus_schema_validation = translateSchema_json2cerberus(schemaObj = json_schema_validation)['Root']['schema']
    cerberus_schema_validation_withDefault = translateSchema_json2cerberus(schemaObj = json_schema_validation, defaultValues = defaultValues_dict)['Root']['schema']

    # Create Validator objects
    validator_schema_validation = cerberus.Validator(cerberus_schema_validation, ignore_none_values = True, purge_unknown = True)
    validator_schema_validation_withDefaults = cerberus.Validator(cerberus_schema_validation_withDefault, ignore_none_values = True, purge_unknown = True)

    # Cerberus: Normalize with Validators
    json_sample_validated = validator_schema_validation.validate(json_sample)  # BEWARE of URL quoting
    json_sample_empty = validator_schema_validation_withDefaults.normalized({'sample': {}})

    if json_sample_validated:
        logging.info(f'JSON Schema Validation of {snakemake.params["sample"]}: Passed.')
    else:
        logging.warning(f'JSON Schema Validation of {snakemake.params["sample"]}: FAILED.')
        with open(snakemake.output['json_out'].replace('.json', '_validationErrors.json'), 'w') as file:
            json.dump(validator_schema_validation.errors,
                      fp = file,
                      indent = 4,
                      sort_keys = False,
                      ensure_ascii = False,
                      allow_nan = False)


def filter_sample(json_sample, json_schema_filter):
    # Cerberus: Create Validator - Convert custom JSON schema to Cerberus schema
    cerberus_schema_filter = translateSchema_json2cerberus(schemaObj = json_schema_filter)['Root']['schema']

    # Create Validator objects
    validator_schema_filter = cerberus.Validator(cerberus_schema_filter, ignore_none_values = True, purge_unknown = True)

    # Cerberus: Normalize with Validators
    json_sample_filtered = validator_schema_filter.normalized(json_sample)  # BEWARE of URL quoting

    if json_sample_filtered == json_sample:
        logging.info(f'JSON Schema Normalization of {snakemake.params["sample"]}: No changes.')
    else:
        logging.warning(f'Schema Normalization of {snakemake.params["sample"]}: Filtered.')

    return json_sample_filtered


# %% Main Function ############################################################

def main():
    debug_switch = False

    global json_sample
    global json_schema_validation
    global json_schema_filter

    # Parse Arguments
    python_or_snakemake(debug = debug_switch)

    # Configure Logging
    logging.basicConfig(filename = snakemake.log['log'],
                        # encoding='utf-8',  # NOTE: enabled since upgrade to python v3.9
                        level = logging.INFO,
                        format = '%(asctime)s %(levelname)-8s %(message)s',
                        datefmt = '%Y-%m-%d %H:%M:%S')
    logging.info("json_filter.py is runnning in DEBUG mode") if debug_switch else None

    # Workflow
    import_data()
    validate_sample(json_sample = json_sample, json_schema_validation = json_schema_validation)
    json_sample_filtered = filter_sample(json_sample = json_sample, json_schema_filter = json_schema_filter)

    # Write Output
    if json_sample_filtered:
        with open(snakemake.output['json_out'], 'w') as file:
            json.dump(json_sample_filtered,
                      fp = file,
                      indent = 4,
                      sort_keys = False,
                      ensure_ascii = False,
                      allow_nan = False)
    else:
        sys.exit(1)

    logging.info(f'JSON Schema procedures for {snakemake.params["sample"]} finished.')
    print(f'JSON Schema procedures for {snakemake.params["sample"]} finished.')


# %% Main Call ################################################################

if __name__ == '__main__':
    main()
