#!/usr/bin/env python3

print('''
parse_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 socket
import re
from datetime import datetime
import json
from numbers import Integral
from copy import deepcopy

# Other Packages
import pandas as pd
from helper_functions import *


# %% Debugging and Logging

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


try:
    snakemake.input[0]
except NameError:
    from parse_json_debug import *
else:
    sys.excepthook = handle_exception  # Install exception handler

logging.basicConfig(filename = snakemake.log['log'],
                    # encoding='utf-8',  # TODO: enable this with upgrade to python v3.9
                    level = logging.INFO,
                    format = '%(asctime)s %(levelname)-8s %(message)s',
                    datefmt = '%Y-%m-%d %H:%M:%S')


# %% Module Set Functions
# region Module Set Functions
def aggregate_versions_and_hashes():
    # %% Module Software Versions & Database Hashes

    ## Define Global Variables
    global version_dict
    global database_hashes

    ## Get Software Versions
    version_report = pd.read_csv(snakemake.input['software_versions'], sep = '\t')
    version_dict = {k: v for (k, v) in dict(zip(version_report.Software, version_report.Version)).items() if not pd.isnull(k)}

    ## Get Database MD5 Hashes
    database_report = pd.read_csv(snakemake.input['db_hashes'], delim_whitespace = True, names = ["Hash", "Path"])
    database_report['File'] = [os.path.basename(x) for x in database_report.Path]
    database_hashes = {k: v for (k, v) in dict(zip(database_report.File, database_report.Hash)).items() if not pd.isnull(k)}


def aggregate_run_metadata(timestamp_list: list = []) -> dict:
    ## Get Input Sample MD5 Hashes, Note: the complexity is due to the copy2scratch strategy in the SLURM branch
    if snakemake.rule == 'parse_json_before_assembly':
        fq1_path = snakemake.input.get('r1', '')
        if os.path.isfile(path = snakemake.input.get('r1', '')):
            fq1_md5 = checksum_md5(filename = snakemake.input.get('r1', ''))
        else:
            fq1_md5 = ''

        fq2_path = snakemake.input.get('r2', '')
        if os.path.isfile(path = snakemake.input.get('r2', '')):
            fq2_md5 = checksum_md5(filename = snakemake.input.get('r2', ''))
        else:
            fq2_md5 = ''

    if snakemake.rule == 'parse_json_after_assembly':
        fq1_path = snakemake.params.get('r1', '')
        if os.path.isfile(path = snakemake.params.get('r1', '')):
            fq1_md5 = checksum_md5(filename = snakemake.params.get('r1', ''))
        else:
            fq1_md5 = ''

        fq2_path = snakemake.params.get('r2', '')
        if os.path.isfile(path = snakemake.params.get('r2', '')):
            fq2_md5 = checksum_md5(filename = snakemake.params.get('r2', ''))
        else:
            fq2_md5 = ''

    ## Replace static container path with host volume path
    config_edit = deepcopy(snakemake.params['config'])
    if snakemake.params['config']['smk_params']['docker']:
        docker_path_replace(nestedObj = config_edit, newPath = snakemake.params['config']['smk_params']['docker'])
        fq1_path = fq1_path.replace('/AQUAMIS/analysis', snakemake.params['config']['smk_params']['docker'])
        fq2_path = fq2_path.replace('/AQUAMIS/analysis', snakemake.params['config']['smk_params']['docker'])

    # %% Module Run Metadata

    json_module_run_metadata = {
        'pipeline'            : "AQUAMIS",
        'timestamp'           : str(max((x for x in timestamp_list if x is not None), default = datetime.now())),
        'sample_description'  : {
            'sample'    : snakemake.params['sample'],
            'read1_file': fq1_path,
            'read2_file': fq2_path,
            'read1_md5' : fq1_md5,
            'read2_md5' : fq2_md5
        },
        'pipeline_metadata'   : {
            'stage'      : snakemake.rule,
            'host_system': socket.gethostname(),
            'config'     : config_edit,
            'software'   : version_dict

        },
        'database_information': {
            'md5_hashes': {**database_hashes}
        }
    }

    json_module_analysis = {
        'run_metadata': [{**json_module_run_metadata}]
    }

    return (json_module_analysis)


def parse_json_before_assembly() -> dict:
    timestamps = list()

    # %% Module FastP

    ## Get Trimmed Fastq MD5 Hashes
    if os.path.isfile(path = snakemake.input.get('trimmed_r1', '')):
        trimmed_R1_md5 = checksum_md5(filename = os.path.join(snakemake.params['config']['workdir'], snakemake.input['trimmed_r1']))
    else:
        trimmed_R1_md5 = ''

    if os.path.isfile(path = snakemake.input.get('trimmed_r2', '')):
        trimmed_R2_md5 = checksum_md5(filename = os.path.join(snakemake.params['config']['workdir'], snakemake.input['trimmed_r2']))
    else:
        trimmed_R2_md5 = ''

    ## FastP report
    timestamps.append(mtime_file(snakemake.input['fastp']))
    with open(snakemake.input['fastp'], 'r') as file:
        fastp_json = json.load(file)

    fastp_json['version'] = version_dict['fastp']
    fastp_json['timestamp'] = str(timestamps[-1])

    ## Select Key/Value pairs for AQUAMIS report
    fastp_json['report'] = {  # TODO: adapter_percentage
        'sample'            : snakemake.params['sample'],
        'total_reads_before': fastp_json['summary']['before_filtering']['total_reads'],
        'total_bases_before': fastp_json['summary']['before_filtering']['total_bases'],
        'total_reads_after' : fastp_json['summary']['after_filtering']['total_reads'],
        'total_bases_after' : fastp_json['summary']['after_filtering']['total_bases'],
        'q20_rate_after'    : fastp_json['summary']['after_filtering']['q20_rate'],
        'q30_rate_after'    : fastp_json['summary']['after_filtering']['q30_rate'],
        'duplication_rate'  : fastp_json['duplication']['rate'],
        'insert_size_peak'  : fastp_json['insert_size']['peak'] if 'insert_size' in fastp_json else None,
        "read1_trimmed_md5" : trimmed_R1_md5,
        "read2_trimmed_md5" : trimmed_R2_md5,
        'link_to_report'    : os.path.join('..', 'trimmed', 'reports', snakemake.params['sample'] + '.html')
    }

    if 'adapter_cutting' in fastp_json:  # conditional due to single-end mode
        adapter_df = pd.DataFrame.from_dict(data = fastp_json['adapter_cutting']['read1_adapter_counts'], orient = 'index', columns = ['counts'])
        adapter_df.index.name = 'adapter'
        fastp_json['adapter_cutting']['read1_adapter_counts'] = adapter_df.reset_index()

        adapter_df = pd.DataFrame.from_dict(data = fastp_json['adapter_cutting']['read2_adapter_counts'], orient = 'index', columns = ['counts'])
        adapter_df.index.name = 'adapter'
        fastp_json['adapter_cutting']['read2_adapter_counts'] = adapter_df.reset_index()

    kmer_df = pd.DataFrame.from_dict(data = fastp_json['read1_before_filtering']['kmer_count'], orient = 'index', columns = ['count'])
    kmer_df.index.name = 'kmer'
    fastp_json['read1_before_filtering']['kmer_count'] = kmer_df.reset_index()

    kmer_df = pd.DataFrame.from_dict(data = fastp_json['read1_after_filtering']['kmer_count'], orient = 'index', columns = ['count'])
    kmer_df.index.name = 'kmer'
    fastp_json['read1_after_filtering']['kmer_count'] = kmer_df.reset_index()

    if "read2_before_filtering" in fastp_json:  # conditional due to single-end mode
        kmer_df = pd.DataFrame.from_dict(data = fastp_json['read2_before_filtering']['kmer_count'], orient = 'index', columns = ['count'])
        kmer_df.index.name = 'kmer'
        fastp_json['read2_before_filtering']['kmer_count'] = kmer_df.reset_index()

        kmer_df = pd.DataFrame.from_dict(data = fastp_json['read2_after_filtering']['kmer_count'], orient = 'index', columns = ['count'])
        kmer_df.index.name = 'kmer'
        fastp_json['read2_after_filtering']['kmer_count'] = kmer_df.reset_index()

    json_module_fastp = {
        'fastp': fastp_json
    }

    # %% Module ConFindr

    try:
        open(snakemake.input['confindr_report'])
    except FileNotFoundError:
        confindr_fail = True
    else:
        confindr_fail = False

    if not confindr_fail:
        timestamps.append(mtime_file(snakemake.input['confindr_report']))
        confindr_report = pd.read_csv(snakemake.input['confindr_report'], sep = ',', dtype = ({0: 'str'}))
        confindr_report_dict = {
            'version'  : version_dict['ConFindr'],
            'timestamp': str(timestamps[-1])
        }

        confindr_report_dict.update(confindr_report.to_dict(orient = 'records')[0])  # only valid if confindr report contains ever only one record
        # ND values in contaminated samples are string, while other samples are integer, combining samples therefore fails. mitigation: convert to None

        if confindr_report_dict['ContamStatus'] is True:
            if ":" in confindr_report_dict['Genus']:
                if confindr_report_dict['NumContamSNVs'] >= 3:  # as defined by max(AQUAMIS_thresholds.json["thresholds"]["all Species"]["NumContamSNVs"]["interval"])
                    confindr_report_dict['ContamStatus'] = "Inter+Intra"
                else:
                    confindr_report_dict['ContamStatus'] = "Inter"
            else:
                confindr_report_dict['ContamStatus'] = "Intra"
        elif confindr_report_dict['ContamStatus'] != 'ND':
            confindr_report_dict['ContamStatus'] = "False"

        if confindr_report_dict['BasesExamined'] is 0:
            confindr_report_dict['ContamStatus'] = None
            confindr_report_dict['NumContamSNVs'] = None

        # Identify database in use
        confindr_log_commands = []
        with open(os.path.join(os.path.dirname(snakemake.input['confindr_report']), "confindr_log.txt"), 'r') as file:
            for line in file:
                if 'Command used' in line:
                    confindr_log_commands.append(line)

        if any('kma' in s for s in confindr_log_commands):
            confindr_log_commands_filtered = [cmd for cmd in confindr_log_commands if 'kma' in cmd]  # + snakemake.params['config']['params']['confindr']['database']
            confindr_report_dict['Database'] = [os.path.basename(re.compile("-t_db (\S*)").findall(cmd)[0]) for cmd in confindr_log_commands_filtered][0].rstrip('_kma')
        else:
            confindr_report_dict['Database'] = 'ND'

        confindr_report_dict = {key: (None if value == 'ND' else value) for (key, value) in confindr_report_dict.items()}

        json_module_confindr = {
            'confindr': confindr_report_dict
        }
    else:
        json_module_confindr = None

    # %% Module Kraken2_Reads

    kraken_reads_report_columns = ['perc', 'tot_all', 'tot_lvl', 'clade_lvl', 'clade_taxid', 'clade_name']
    timestamps.append(mtime_file(snakemake.input['kraken_reads_report']))
    kraken_reads_report = pd.read_csv(snakemake.input['kraken_reads_report'], sep = '\t', names = kraken_reads_report_columns, nrows = 50)
    kraken_reads_report['clade_name'] = kraken_reads_report['clade_name'].str.strip()

    bracken_report_columns = ['name', 'taxonomy_id', 'taxonomy_lvl', 'kraken_assigned_reads', 'added_reads', 'new_est_reads', 'fraction_total_reads']
    bracken_species_report = pd.read_csv(snakemake.input['bracken_species_report'], sep = '\t', names = bracken_report_columns, nrows = 50)
    bracken_species_report['name'] = bracken_species_report['name'].str.strip()
    bracken_genus_report = pd.read_csv(snakemake.input['bracken_genus_report'], sep = '\t', names = bracken_report_columns, nrows = 50)
    bracken_genus_report['name'] = bracken_genus_report['name'].str.strip()

    json_module_kraken2_reads = {
        'kraken2': {
            'read_based': {
                'version'                : version_dict['Kraken'],
                'database_md5'           : database_hashes['hash.k2d'],
                'timestamp'              : str(timestamps[-1]),
                'kraken_report'          : kraken_reads_report,
                'bracken_summary_species': bracken_species_report,
                'bracken_summary_genus'  : bracken_genus_report
            }
        }
    }

    # %% JSON Module Assembly - BEFORE Assembly

    json_module_analysis = aggregate_run_metadata(timestamp_list = timestamps)

    json_assembly = {
        'sample'   : {
            **json_module_analysis
        },
        'pipelines': {
            **json_module_fastp,
            **json_module_confindr,
            **json_module_kraken2_reads,
            'aquamis': {**json_module_analysis}
        }
    }

    return (json_assembly)


def parse_json_after_assembly() -> dict:
    timestamps = list()

    # %% Module Shovill

    assembly_fasta_header_lines = []
    with open(snakemake.input['assembly'], 'r') as file:
        for line in file:
            if '>' in line:
                assembly_fasta_header_lines.append(line.strip('\n').split(' '))

    shovill_columns = ['>', 'len', 'cov', 'corr', 'origname', 'sw', 'date']
    assembly_fasta_header = pd.DataFrame(assembly_fasta_header_lines, columns = shovill_columns)
    assembly_fasta_header = pd.DataFrame({name: assembly_fasta_header[name].replace(rf'{name}=*', '', regex = True) for name in assembly_fasta_header.columns.values})
    assembly_fasta_header = assembly_fasta_header.apply(pd.to_numeric, errors = 'ignore')
    assembly_fasta_header.rename(columns = {'>': 'id'}, inplace = True)

    with open(snakemake.input['assembly_log'], 'r') as file:
        shovill_log = file.readlines()

    assembly_downsampling_factor = 1.0
    flash_total = None
    flash_combined = None
    flash_uncombined = None

    for line in shovill_log:
        if '[shovill]' in line and 'Subsampling reads by factor' in line:
            assembly_downsampling_factor = re.search('0\.\d+|$', line).group()
        if '[FLASH]' in line and 'Total pairs' in line:
            flash_total = re.search('\d+|$', line).group()
        if '[FLASH]' in line and 'Combined pairs' in line:
            flash_combined = re.search('\d+|$', line).group()
        if '[FLASH]' in line and 'Uncombined pairs' in line:
            flash_uncombined = re.search('\d+|$', line).group()

    fasta_file = os.path.join(snakemake.params['config']['workdir'], "Assembly", "assembly", snakemake.params['sample'] + ".fasta")
    timestamps.append(mtime_file(fasta_file))
    fasta_md5 = checksum_md5(filename = fasta_file)

    json_module_shovill = {
        'shovill': {
            'version'              : version_dict['shovill'],
            'timestamp'            : str(timestamps[-1]),
            'fasta_file'           : fasta_file,
            'fasta_md5'            : fasta_md5,
            'downsampling_factor'  : float(assembly_downsampling_factor),
            'flash'                : {
                'total_reads'     : flash_total,
                'combined reads'  : flash_combined,
                'uncombined_reads': flash_uncombined
            },
            'assembly_fasta_header': assembly_fasta_header,
        }
    }

    # %% Module Samtools

    def extractSamStats(samstats, index):
        extract = [x for x in samstats if re.search(rf'^{index}', x) != None]
        extract = [re.sub(rf'^{index}\t|\n', '', x) for x in extract]
        extract = [re.sub(':\t|\t# ', '\t', x) for x in extract]
        df = pd.DataFrame([x.split('\t') for x in extract])
        return df

    timestamps.append(mtime_file(snakemake.input['samstats']))

    with open(snakemake.input['samstats'], 'r') as file:
        samstats = file.readlines()

    samstats_index = ['SN', 'FFQ', 'LFQ', 'GCF', 'GCL', 'GCC', 'FBC', 'FTC', 'LBC', 'LTC', 'IS', 'RL', 'FRL', 'LRL', 'ID', 'IC', 'COV', 'GCD']

    samstats_summary = extractSamStats(samstats = samstats, index = 'SN').set_index(0)
    samstats_summary.columns = ['value', 'comment']
    samstats_summary_dict = pd.DataFrame(samstats_summary.value).T.apply(pd.to_numeric, errors = 'ignore').to_dict(orient = 'records')[0]
    samstats_summary_dict['comments'] = {k: v for (k, v) in dict(zip(samstats_summary.index, samstats_summary.comment)).items() if not pd.isnull(v)}
    samstats_summary_dict['reads mapped count'] = int(pd.read_csv(snakemake.input['samstats_view_count'], sep = '\t', header = None).iloc[0, 0])

    samstats_insertsize = extractSamStats(samstats = samstats, index = 'IS')
    # samstats_insertsize_names = ['insert size', 'pairs total', 'inward oriented pairs', 'outward oriented pairs', 'other orientation pairs']  # this is the official notation from the manual at http://www.htslib.org/doc/samtools-stats.html
    samstats_insertsize_names = ['insert_size', 'count', 'inward_count', 'outward_count', 'other_pairs']
    samstats_insertsize = pd.DataFrame(0, index = np.arange(1), columns = samstats_insertsize_names) if samstats_insertsize.empty else samstats_insertsize
    samstats_insertsize.columns = samstats_insertsize_names
    samstats_insertsize = samstats_insertsize.apply(pd.to_numeric)

    samstats_coverage = extractSamStats(samstats = samstats, index = 'COV')
    samstats_coverage.columns = ['coverage_interval', 'coverage_bin', 'frequency']
    samstats_coverage[['coverage_bin', 'frequency']] = samstats_coverage[['coverage_bin', 'frequency']].apply(pd.to_numeric)
    # samstats_coverage_dict = samstats_coverage[['coverage_bin', 'frequency']].to_dict()  # less overhead, but does not parse in JSONschema.net

    json_module_samstats = {
        'samstats': {
            'version'    : version_dict['samtools'],
            'timestamp'  : str(timestamps[-1]),
            'summary'    : samstats_summary_dict,
            'insert_size': samstats_insertsize,
            'coverage'   : samstats_coverage.drop(columns = 'coverage_interval')  # TODO: Decission: If coverage_interval is always '[coverage_bin-coverage_bin]' then use samstats_coverage_dict
        }
    }

    # %% Module Kraken2_Contigs

    kraken_contigs_report_columns = ['contig', 'taxid', 'length', 'lineage']
    kraken_taxon_ranks = {'K': 'kingdom',
                          'P': 'phylum',
                          'C': 'class',
                          'O': 'order',
                          'F': 'family',
                          'G': 'genus',
                          'S': 'species'
                          }

    kraken_taxon_unclassified = 'k__unclassified|p__unclassified|c__unclassified|o__unclassified|f__unclassified|g__unclassified|s__unclassified'
    timestamps.append(mtime_file(snakemake.input['kraken_contigs_report']))
    kraken_contigs_report = pd.read_csv(snakemake.input['kraken_contigs_report'], sep = '\t', names = kraken_contigs_report_columns, skip_blank_lines = True)
    kraken_contigs_report.loc[pd.isnull(kraken_contigs_report['lineage']), 'lineage'] = kraken_taxon_unclassified  # BUGFIX: if result of kraken is unclassified (U in first column of .kraken file) taxonkit produces no result for a contig, there is a blank line
    lineage = pd.DataFrame(kraken_contigs_report['lineage'].str.split('|').values.tolist(),
                           columns = kraken_taxon_ranks.values()).replace(regex = '[kpcofgs]__', value = '')
    kraken_contigs_report = kraken_contigs_report.drop(labels = 'lineage', axis = 1).join(lineage)

    taxon_level = kraken_taxon_ranks['S']
    taxon_sum = {taxon: kraken_contigs_report.loc[kraken_contigs_report[taxon_level] == taxon, 'length'].sum() for taxon in kraken_contigs_report[taxon_level].unique()}
    taxon_fraction = {taxon: (length / kraken_contigs_report['length'].sum()) for (taxon, length) in taxon_sum.items()}
    taxonkit_species_summary = pd.DataFrame({taxon_level + '_cumlength': pd.Series(taxon_sum), taxon_level + '_fraction': pd.Series(taxon_fraction)}).sort_values(by = [taxon_level + '_cumlength'], ascending = False)
    taxonkit_species_summary.index.name = taxon_level
    taxonkit_species_summary.reset_index(inplace = True)
    taxonkit_species_summary.insert(loc = 0, column = 'sample', value = snakemake.params['sample'])

    taxon_level = kraken_taxon_ranks['G']
    taxon_sum = {taxon: kraken_contigs_report.loc[kraken_contigs_report[taxon_level] == taxon, 'length'].sum() for taxon in kraken_contigs_report[taxon_level].unique()}
    taxon_fraction = {taxon: (length / kraken_contigs_report['length'].sum()) for (taxon, length) in taxon_sum.items()}
    taxonkit_genus_summary = pd.DataFrame({taxon_level + '_cumlength': pd.Series(taxon_sum), taxon_level + '_fraction': pd.Series(taxon_fraction)}).sort_values(by = [taxon_level + '_cumlength'], ascending = False)
    taxonkit_genus_summary.index.name = taxon_level
    taxonkit_genus_summary.reset_index(inplace = True)
    taxonkit_genus_summary.insert(loc = 0, column = 'sample', value = snakemake.params['sample'])

    json_module_kraken2_contigs = {
        'kraken2': {
            'contig_based': {
                'version'               : version_dict['Kraken'],
                'database_md5'          : database_hashes['hash.k2d'],
                'timestamp'             : str(timestamps[-1]),
                'kraken_report'         : kraken_contigs_report,
                'kraken_summary_species': taxonkit_species_summary,
                'kraken_summary_genus'  : taxonkit_genus_summary,
            }
        }
    }

    # %% Module Mash

    mash_distance_columns = ['reference', 'query', 'distance', 'pvalue', 'shared_hashes', 'total_hashes']
    timestamps.append(mtime_file(snakemake.input['mash_distance']))
    mash_distance = pd.read_csv(snakemake.input['mash_distance'], sep = '\t', names = mash_distance_columns, nrows = 1)
    mash_distance = mash_distance.to_dict(orient = 'records')[0]

    with open(snakemake.input['reference'], 'r') as file:
        mash_reference = re.sub(r'^>(.*)\n', '\\1', file.readline())

    json_module_mash = {
        'mash': {
            'version'         : version_dict['mash'],
            'database_md5'    : database_hashes['mashDB.msh'],
            'timestamp'       : str(timestamps[-1]),
            'reference_header': mash_reference,
            'reference'       : mash_distance.pop('reference'),
            **mash_distance
        }
    }

    # %% Module QUAST

    timestamps.append(mtime_file(snakemake.input['quast_report']))
    quast_report = pd.read_csv(snakemake.input['quast_report'], sep = '\t')
    quast_report = quast_report.set_index('Assembly').T
    quast_report = quast_report.apply(pd.to_numeric, errors = 'ignore').to_dict(orient = 'records')[0]

    quast_misassemblies = pd.read_csv(snakemake.input['quast_misassemblies'], sep = '\t')
    quast_misassemblies.columns = quast_misassemblies.columns.str.strip()
    quast_misassemblies = quast_misassemblies.to_dict(orient = 'records')[0]

    quast_unaligned = pd.read_csv(snakemake.input['quast_unaligned'], sep = '\t')
    quast_unaligned = quast_unaligned.set_index('Assembly').T.to_dict(orient = 'records')[0]

    ## Parse QUAST genome_info file
    with open(snakemake.input['quast_genome_info'], 'r') as file:
        quast_genome_info_lines = file.readlines()

    # TODO: bugfix needed for multi-fasta reference files downloaded from NCBI assembly
    quast_genome_info_lines = [x.strip() for x in quast_genome_info_lines]
    quast_genome_info_lines[0] = quast_genome_info_lines[0].replace('"', '')
    quast_genome_info_lines[1] = quast_genome_info_lines[1] + ' ' + re.sub(r'^(\S*)\s.*', '\\1', quast_genome_info_lines[2])
    quast_genome_info_lines[2] = ['reference {0}'.format(i) for i in re.sub(r'.*\((.*)\).*', '\\1', quast_genome_info_lines[2]).split(', ')]
    quast_genome_info_tableindex_1st = [s.split('|') for s in quast_genome_info_lines if 'assembly' in s][0]
    quast_genome_info_tableindex_2nd = [s.split('|') for s in quast_genome_info_lines if 'fraction' in s][0]
    quast_genome_info_tableindex = [(quast_genome_info_tableindex_1st[n].strip() + " " + quast_genome_info_tableindex_2nd[n].strip()) for n, val in enumerate(quast_genome_info_tableindex_1st)]
    quast_genome_name_trunc = snakemake.params['sample'].replace('.', '-')
    quast_genome_name_trunc = quast_genome_name_trunc[:24] if len(quast_genome_name_trunc) > 24 else quast_genome_name_trunc
    quast_genome_info_tabledata = [s.split('|') for s in quast_genome_info_lines if quast_genome_name_trunc in s]
    quast_genome_info = dict(x.split(': ') for x in quast_genome_info_lines if ':' in x)
    quast_genome_info.update(dict(x.split(': ') for x in quast_genome_info_lines[2]))
    quast_genome_info_table = dict(zip([s.strip() for s in quast_genome_info_tableindex], [s.strip() for s in quast_genome_info_tabledata[0]]))
    quast_genome_info.update({k: v for (k, v) in quast_genome_info_table.items() if k != ''})
    quast_genome_info = pd.DataFrame.from_dict(
        quast_genome_info, orient = 'index').replace(
        ' bp', '', regex = True).T.apply(
        pd.to_numeric, errors = 'ignore').to_dict(orient = 'records')[0]

    json_module_quast = {
        'quast': {
            'version'      : version_dict['QUAST'],
            'timestamp'    : str(timestamps[-1]),
            'report'       : quast_report,
            'genome_info'  : quast_genome_info,
            'misassemblies': quast_misassemblies,
            'unaligned'    : quast_unaligned
        }
    }

    # %% Module BUSCO

    busco_file = os.path.join(os.path.dirname(snakemake.params['busco']),
                              os.path.splitext(os.path.basename(snakemake.params['busco']))[0].replace('.', '-') +
                              os.path.splitext(os.path.basename(snakemake.params['busco']))[1])

    try:
        open(busco_file)
    except FileNotFoundError:
        busco_fail = True
    else:
        busco_fail = False

    if not busco_fail:
        timestamps.append(mtime_file(busco_file))
        with open(busco_file) as file:
            busco_summary = file.readlines()

        busco_summary = [x.strip() for x in busco_summary]

        busco_summary_notes = list(filter(re.compile('[#].*').match, busco_summary))  # Read Note
        busco_summary_notes = [x for x in busco_summary_notes if ':' in x]
        busco_summary_notes_dict = dict(x.split(': ', maxsplit = 1) for x in busco_summary_notes)

        busco_summary_results = list(filter(re.compile('[^#].*').match, busco_summary))  # Read Note
        busco_summary_results_stats = re.split('[,\[\]]', busco_summary_results[0])
        busco_summary_results_stats = [x for x in busco_summary_results_stats if x]
        busco_summary_results_stats = dict(x.split(':') for x in busco_summary_results_stats)

        busco_summary_results_dict = dict(reversed(x.split('\t')) for x in ';'.join(busco_summary_results[1:]).split(';'))
        busco_summary_results_dict = {value: int(key) for value, key in busco_summary_results_dict.items()}

        json_module_busco = {
            'busco': {
                'version'  : busco_summary_notes_dict['# BUSCO version is'],
                'timestamp': str(timestamps[-1]),
                'command'  : busco_summary_notes_dict['# To reproduce this run'],
                'database' : busco_summary_notes_dict['# The lineage dataset is'],
                'results'  : {**busco_summary_results_dict, **busco_summary_results_stats}
            }
        }
    else:
        json_module_busco = None

    # %% Module MLST

    timestamps.append(mtime_file(snakemake.input['mlst']))
    with open(snakemake.input['mlst'], 'r') as file:
        mlst_json = json.load(file)[0]  # prerequisite: only single-sample-json loaded

    if mlst_json['scheme'] is not "-":
        mlst_json['loci_total'] = len(mlst_json['alleles'])
        mlst_json['loci_missing'] = len([v for v in mlst_json['alleles'].values() if v is '-'])
        mlst_json['multi_allele_loci'] = sum([len(re.findall(",", locus)) for locus in mlst_json['alleles'].values()])
        ## reorder alleles by locus name
        mlst_json_alleles_df = pd.DataFrame.from_dict(data = mlst_json['alleles'], orient = 'index', columns = ['allele_no'])
        mlst_json_alleles_df.index.name = 'locus'
        mlst_json['alleles'] = mlst_json_alleles_df.sort_values(by = 'locus', axis = 0).reset_index().to_json(orient = 'records')
    else:
        mlst_json['loci_total'] = None
        mlst_json['loci_missing'] = None
        mlst_json['multi_allele_loci'] = None

    ## stabilize key order
    mlst_json_sorted = json.loads(json.dumps(obj = mlst_json, sort_keys = True))

    json_module_mlst = {
        'mlst': {
            'version'  : version_dict['mlst'],
            'timestamp': str(timestamps[-1]),
            'id'       : mlst_json_sorted.pop('id'),
            'filename' : mlst_json_sorted.pop('filename'),
            'scheme'   : mlst_json_sorted.pop('scheme'),
            **mlst_json_sorted
        }
    }

    # %% Module AQUAMIS

    aquamis_report = {'sample': snakemake.params['sample']}
    aquamis_report['reference'] = json_module_mash['mash']['reference_header']
    aquamis_report['assembly_coverageDepth'] = sum(samstats_coverage['coverage_bin'] * samstats_coverage['frequency']) / sum(samstats_coverage['frequency'])
    aquamis_report['reads_total'] = json_module_samstats['samstats']['summary']['raw total sequences']
    aquamis_report['reads_mapped'] = json_module_samstats['samstats']['summary']['reads mapped']
    aquamis_report['fraction_reads_mapped'] = aquamis_report['reads_mapped'] / aquamis_report['reads_total']
    aquamis_report['coverage_rawDepth'] = json_module_samstats['samstats']['summary']['reads mapped count'] / sum(json_module_shovill['shovill']['assembly_fasta_header']['len'])
    aquamis_report['base_count'] = json_module_samstats['samstats']['summary']['reads mapped count']
    aquamis_report['count_circular_contigs'] = None  # TODO: adapt to shovill or remove

    try:
        aquamis_report['busco_single'] = float(json_module_busco['busco']['results']['S'].strip('%'))  # single copy orthologs
        aquamis_report['busco_duplicates'] = float(json_module_busco['busco']['results']['D'].strip('%'))  # duplicated orthologs
    except TypeError:
        aquamis_report['busco_single'] = None
        aquamis_report['busco_duplicates'] = None

    aquamis_report['genes_full'] = int(re.findall('\d+', json_module_quast['quast']['report']['# genomic features'])[0])

    try:
        aquamis_report['genes_partial'] = int(re.findall('\d+', json_module_quast['quast']['report']['# genomic features'])[1])
    except IndexError:
        aquamis_report['genes_partial'] = ''

    aquamis_report['fraction_genes_recovered'] = (aquamis_report['genes_full'] + aquamis_report['genes_partial']) / json_module_quast['quast']['genome_info']['Genomic features of type ANY loaded']
    # aquamis_report['ref_length'] = json_module_quast['quast']['report']['Reference length']  # Calculated in former snakemake rule but not used in report
    aquamis_report['version'] = snakemake.params['config']['version']  # duplicate in config, but here it conforms to module structure
    aquamis_report['timestamp'] = str(max((x for x in timestamps if x is not None), default = datetime.now()))

    json_module_aquamis = {
        'aquamis': {
            'report': aquamis_report
        }
    }

    # %% Module Sample Summary

    # TODO: this summary serves as a EFSA/BeONE Minimal Subset template, revise appropriate keys.

    # collect result keys for a Minimal Subset
    subset = {
        "aquamis": [
            'sample',
            'reference',
            'assembly_coverageDepth',
            'reads_total',
            'reads_mapped',
            'fraction_reads_mapped',
            'coverage_rawDepth',
            'base_count',
            'busco_single',
            'busco_duplicates',
            'genes_full',
            'genes_partial',
            'fraction_genes_recovered'
        ],
        "quast"  : [
            '# contigs (>= 0 bp)',
            '# contigs (>= 1000 bp)',
            '# contigs (>= 10000 bp)',
            'Total length (>= 0 bp)',
            'Total length (>= 1000 bp)',
            'Total length (>= 10000 bp)',
            'Reference length',
            'GC (%)',
            'Reference GC (%)',
            'N50',
            'Genome fraction (%)',
            "# N's per 100 kbp",
            'Complete BUSCO (%)',
            'Partial BUSCO (%)',
            'Total aligned length'
        ],
        "kraken2": [
            'species',
            'species_cumlength',
            'species_fraction'
        ],
        "mlst"   : [
            'sequence_type'
        ]
    }

    _kraken_df = json_module_kraken2_contigs['kraken2']['contig_based']['kraken_summary_species']

    # subset json_modules
    json_module_summary_subset = {
        **{key: value for key, value in json_module_aquamis['aquamis']['report'].items() if key in subset["aquamis"]},
        **{key: value for key, value in json_module_quast['quast']['report'].items() if key in subset["quast"]},
        **_kraken_df.loc[0, _kraken_df.columns.intersection(subset["kraken2"])].to_dict(),
        **{key: value for key, value in json_module_mlst['mlst'].items() if key in subset["mlst"]}
    }

    # convert numpy int64 to int
    json_module_summary_subset_int = {key: (int(value) if isinstance(value, Integral) else value) for key, value in json_module_summary_subset.items()}

    # round floats and convert numpy float64 to float
    json_module_summary_subset_round = {key: (float(round(value, 5)) if isinstance(value, float) else value) for key, value in json_module_summary_subset_int.items()}

    json_module_summary = {
        'summary': {
            **json_module_summary_subset_round
        }
    }

    # %% JSON Module Assembly - AFTER Assembly

    with open(snakemake.input['json_in'], 'r') as file:
        json_sample = json.load(fp = file, object_pairs_hook = unquote_hook)

    json_module_analysis = aggregate_run_metadata(timestamp_list = timestamps)
    json_module_analysis['run_metadata'][0]['sample_description'].update({
        'fasta_file': fasta_file,
        'fasta_md5' : fasta_md5
    })

    try:
        json_sample['sample']['run_metadata'].append(json_module_analysis['run_metadata'][0])
        json_sample['pipelines']['aquamis']['run_metadata'].append(json_module_analysis['run_metadata'][0])
    except KeyError:
        pass

    json_assembly = {
        'sample'   : {
            **json_module_summary
        },
        'pipelines': {
            **json_module_shovill,
            **json_module_samstats,
            **json_module_kraken2_contigs,
            **json_module_mash,
            **json_module_quast,
            **json_module_mlst,
            **json_module_aquamis
        }
    }

    try:
        json_assembly['pipelines'].update(json_module_busco)
    except TypeError:
        logging.warning('BUSCO result aggregation into JSON failed.')

    json_assembly_merge = merge_dict(original = json_sample, patch = json_assembly)
    return (json_assembly_merge)


def parse_json_after_qc() -> dict:
    timestamps = list()

    # %% Module QC Assessment

    with open(snakemake.params['config']['params']['qc']['thresholds'], 'r') as file:
        json_threshold_metadata_nested = json.load(fp = file, object_pairs_hook = unquote_hook)['metadata']

    ## Get Database MD5 Hashes
    database_report = pd.read_csv(snakemake.input['db_hashes'], delim_whitespace = True, names = ["Hash", "Path"])
    database_report['File'] = [os.path.basename(x) for x in database_report.Path]
    database_hashes = {k: v for (k, v) in dict(zip(database_report.File, database_report.Hash)).items() if not pd.isnull(k)}

    json_threshold_metadata = {k: v[0] for (k, v) in json_threshold_metadata_nested.items()}
    json_threshold_metadata.pop('user', None)
    json_threshold_metadata['database_md5'] = database_hashes['AQUAMIS_thresholds.json']
    timestamps.append(mtime_file(snakemake.params['tsv_qc_vote']))  # TODO: stabilize qc timestamp
    qc_report = pd.read_csv(snakemake.params['tsv_qc_vote'], sep = '\t')
    qc_report_dict = qc_report[qc_report['sample'].eq(snakemake.params['sample'])].to_dict(orient = 'records')[0]

    json_module_qc = {
        'qc_assessment': {
            'thresholds': json_threshold_metadata,
            'timestamp' : str(timestamps[-1]),
            **qc_report_dict,
        }
    }

    # %% JSON Module Assembly - AFTER QC

    with open(snakemake.input['json_in'], 'r') as file:
        json_sample = json.load(fp = file, object_pairs_hook = unquote_hook)

    json_assembly = {
        'sample'   : {
            'qc_assessment': {
                **{k: v for (k, v) in json_module_qc['qc_assessment'].items() if re.compile('^qc_').match(k)}  # Minimal Subset: distill counts and overall vote
            }
        },
        'pipelines': {
            'aquamis': {
                **json_module_qc
            }
        }
    }

    json_assembly = merge_dict(original = json_sample, patch = json_assembly)
    return (json_assembly)


# endregion
# %% JSON Aggregation Module Sets

logging.info('Starting Module Set: ' + snakemake.rule)
aggregate_versions_and_hashes()

if snakemake.rule == 'parse_json_before_assembly':
    json_assembly = parse_json_before_assembly()

if snakemake.rule == 'parse_json_after_assembly':
    json_assembly = parse_json_after_assembly()

if snakemake.rule == 'parse_json_after_qc':
    json_assembly = parse_json_after_qc()

# %% Export JSON Output

with open(snakemake.output['json_out'], 'w') as file:
    json.dump(
        obj = strictJsonObj(
            nestedObj = json.loads(
                s = nestedDict_to_json(
                    nestedObj = json_assembly)),
            uriCompliant = True),
        fp = file,
        indent = 4,
        sort_keys = False,
        ensure_ascii = False,
        allow_nan = False)  # Improvement: Use .encode('utf-8') in combo with ensure_ascii=False if None values are omitted instead of coerced to None

logging.info('JSON Assembly for ' + snakemake.params['sample'] + ' finished.')
