# Path: kleborate/__main__.py
# def main():
# def parse_arguments():
# def check_inputs_and_programs(args):
# def get_blast_version():
# def get_output_headers(args, data_folder):
# def get_virulence_score(yersiniabactin, colibactin, aerobactin):
# def get_resistance_score(res_headers, res_hits):
# def get_resistance_class_count(res_headers, res_hits):
# def get_resistance_gene_count(res_headers, res_hits):
# def get_data_path():
# def get_chromosome_mlst_header():
# def get_ybt_mlst_header():
# def get_clb_mlst_header():
# def get_iuc_mlst_header():
# def get_iro_mlst_header():
# def get_rmp_mlst_header():
# def gunzip_contigs_if_necessary(contigs, temp_dir):
# def decompress_file(in_file, out_file):
# def get_chromosome_mlst_results(data_folder, contigs, kp_complex, args):
# def get_kp_subspecies_based_on_st(chr_st):
# def get_virulence_cluster_results(data_folder, contigs, alleles_fasta, profiles_txt,
#                                   vir_name, vir_st_name, unknown_group_name, min_gene_count,
#                                   header_function, args):
# def get_all_virulence_results(data_folder, contigs, args):
# def get_ybt_mlst_results(data_folder, contigs, args):
# def get_clb_mlst_results(data_folder, contigs, args):
# def get_iuc_mlst_results(data_folder, contigs, args):
# def get_iro_mlst_results(data_folder, contigs, args):
# def get_rmp_mlst_results(data_folder, contigs, args):
# def get_rmpa2_results(data_folder, contigs, args):
# def get_wzi_and_k_locus_results(data_folder, contigs, args):
# def get_resistance_results(data_folder, contigs, args, res_headers, kp_complex):
# def get_summary_results(results, res_headers, args):
# def output_headers(stdout_header, full_header, outfile):
# def output_results(stdout_header, full_header, outfile, results):
# def get_strain_name(full_path):
# def rebuild_blast_indices(data_dir):


import collections
import tempfile
import os
import unittest
from kleborate.__main__ import get_output_headers, get_strain_name, get_contig_stat_results, \
    get_species_results, is_kp_complex, get_chromosome_mlst_results, get_ybt_mlst_results, \
    get_clb_mlst_results, get_iuc_mlst_results, get_iro_mlst_results, get_rmp_mlst_results,\
    get_rmpa2_results, get_wzi_and_k_locus_results, get_resistance_results, get_summary_results, \
    gunzip_contigs_if_necessary
    return dict(zip(headers, results))


class TestGenomes(unittest.TestCase):

    def setUp(self):
        self.data_dir = 'test/test_genomes/data'
        Args = collections.namedtuple('Args', ['resistance', 'kaptive_k', 'kaptive_o',
                                               'min_coverage', 'min_identity'])
        self.args = Args(resistance=True, kaptive_k=False, kaptive_o=False,
                         min_coverage=80.0, min_identity=90.0)
        _, _, self.res_headers = get_output_headers(self.args, self.data_dir)

    def one_genome_test(self, filename):
        correct_results = load_correct_results(filename)
        with tempfile.TemporaryDirectory() as tmp_dir:
            contigs = gunzip_contigs_if_necessary(filename, tmp_dir)
            test_results = self.get_all_results(contigs)
        for result, value in correct_results.items():
            if value != '':
                assert test_results[result] == value

    def get_all_results(self, contigs):
        results = {'strain': get_strain_name(contigs)}
        results.update(get_species_results(contigs, self.data_dir))
        kp_complex = is_kp_complex(results)
        results.update(get_contig_stat_results(contigs, is_kp_complex))
        results.update(get_chromosome_mlst_results(self.data_dir, contigs, kp_complex, self.args))
        results.update(get_ybt_mlst_results(self.data_dir, contigs, self.args))
        results.update(get_clb_mlst_results(self.data_dir, contigs, self.args))

# Path: kleborate/__main__.py
# def main():
# def parse_arguments():
# def check_inputs_and_programs(args):
# def get_blast_version():
# def get_output_headers(args, data_folder):
# def get_virulence_score(yersiniabactin, colibactin, aerobactin):
# def get_resistance_score(res_headers, res_hits):
# def get_resistance_class_count(res_headers, res_hits):
# def get_resistance_gene_count(res_headers, res_hits):
# def get_data_path():
# def get_chromosome_mlst_header():
# def get_ybt_mlst_header():
# def get_clb_mlst_header():
# def get_iuc_mlst_header():
# def get_iro_mlst_header():
# def get_rmp_mlst_header():
# def gunzip_contigs_if_necessary(contigs, temp_dir):
# def decompress_file(in_file, out_file):
# def get_chromosome_mlst_results(data_folder, contigs, kp_complex, args):
# def get_kp_subspecies_based_on_st(chr_st):
# def get_virulence_cluster_results(data_folder, contigs, alleles_fasta, profiles_txt,
#                                   vir_name, vir_st_name, unknown_group_name, min_gene_count,
#                                   header_function, args):
# def get_all_virulence_results(data_folder, contigs, args):
# def get_ybt_mlst_results(data_folder, contigs, args):
# def get_clb_mlst_results(data_folder, contigs, args):
# def get_iuc_mlst_results(data_folder, contigs, args):
# def get_iro_mlst_results(data_folder, contigs, args):
# def get_rmp_mlst_results(data_folder, contigs, args):
# def get_rmpa2_results(data_folder, contigs, args):
# def get_wzi_and_k_locus_results(data_folder, contigs, args):
# def get_resistance_results(data_folder, contigs, args, res_headers, kp_complex):
# def get_summary_results(results, res_headers, args):
# def output_headers(stdout_header, full_header, outfile):
# def output_results(stdout_header, full_header, outfile, results):
# def get_strain_name(full_path):
# def rebuild_blast_indices(data_dir):


# Path: test/_test_genomes.py
import collections
import tempfile
import os
import unittest
from kleborate.__main__ import get_output_headers, get_strain_name, get_contig_stat_results, \
    get_species_results, is_kp_complex, get_chromosome_mlst_results, get_ybt_mlst_results, \
    get_clb_mlst_results, get_iuc_mlst_results, get_iro_mlst_results, get_rmp_mlst_results,\
    get_rmpa2_results, get_wzi_and_k_locus_results, get_resistance_results, get_summary_results, \
    gunzip_contigs_if_necessary

    return dict(zip(headers, results))


class TestGenomes(unittest.TestCase):

    def setUp(self):
        self.data_dir = 'test/test_genomes/data'
        Args = collections.namedtuple('Args', ['resistance', 'kaptive_k', 'kaptive_o',
                                               'min_coverage', 'min_identity'])
        self.args = Args(resistance=True, kaptive_k=False, kaptive_o=False,
                         min_coverage=80.0, min_identity=90.0)
        _, _, self.res_headers = get_output_headers(self.args, self.data_dir)

    def one_genome_test(self, filename):
        correct_results = load_correct_results(filename)
        with tempfile.TemporaryDirectory() as tmp_dir:
            contigs = gunzip_contigs_if_necessary(filename, tmp_dir)
            test_results = self.get_all_results(contigs)
        for result, value in correct_results.items():
            if value != '':
                assert test_results[result] == value

    def get_all_results(self, contigs):
        results = {'strain': get_strain_name(contigs)}
        results.update(get_species_results(contigs, self.data_dir))
        kp_complex = is_kp_complex(results)
        results.update(get_contig_stat_results(contigs, is_kp_complex))
        results.update(get_chromosome_mlst_results(self.data_dir, contigs, kp_complex, self.args))
        results.update(get_ybt_mlst_results(self.data_dir, contigs, self.args))
        results.update(get_clb_mlst_results(self.data_dir, contigs, self.args))

        results.update(get_iuc_mlst_results(self.data_dir, contigs, self.args))