# Path: kleborate/__main__.py
# def get_output_headers(args, data_folder):
#     """
#     There are two levels of output:
#       * stdout is simpler and displayed to the console
#       * full contains more and is saved to file
#     This function returns headers for both. It also returns the resistance headers in a separate
#     list, as they are used to total up some resistance summaries.
#     """
#     stdout_header = ['strain', 'species']
#     full_header = ['strain', 'species', 'species_match']
#     stdout_header += ['ST', 'virulence_score']
#     full_header += ['contig_count', 'N50', 'largest_contig', 'total_size', 'ambiguous_bases',
#                     'QC_warnings', 'ST', 'virulence_score']
# 
#     if args.resistance:
#         stdout_header.append('resistance_score')
#         full_header.append('resistance_score')
#         full_header.append('num_resistance_classes')
#         full_header.append('num_resistance_genes')
# 
#     other_columns = ['Yersiniabactin', 'YbST',
#                      'Colibactin', 'CbST',
#                      'Aerobactin', 'AbST',
#                      'Salmochelin', 'SmST',
#                      'RmpADC', 'RmST',
#                      'rmpA2',
#                      'wzi', 'K_locus']
#     stdout_header += other_columns
#     full_header += other_columns
# 
#     if args.kaptive_k:
#         stdout_header.append('K_locus_confidence')
#         full_header.append('K_type')
#         full_header.append('K_locus_problems')
#         full_header.append('K_locus_confidence')
#         full_header.append('K_locus_identity')
#         full_header.append('K_locus_missing_genes')
# 
#     if args.kaptive_o:
#         stdout_header.append('O_locus')
#         stdout_header.append('O_locus_confidence')
#         full_header.append('O_locus')
#         full_header.append('O_type')
#         full_header.append('O_locus_problems')
#         full_header.append('O_locus_confidence')
#         full_header.append('O_locus_identity')
#         full_header.append('O_locus_missing_genes')
# 
#     if args.resistance:
#         gene_info, res_classes, bla_classes = \
#             read_class_file(data_folder + '/CARD_AMR_clustered.csv')
#         res_headers = get_res_headers(res_classes, bla_classes)
#         res_headers += ['truncated_resistance_hits', 'spurious_resistance_hits']
#         stdout_header += res_headers
#         full_header += res_headers
#     else:
#         res_headers = []
# 
#     full_header.append('Chr_ST')
#     full_header += get_chromosome_mlst_header()
#     full_header += get_ybt_mlst_header()
#     full_header += get_clb_mlst_header()
#     full_header += get_iuc_mlst_header()
#     full_header += get_iro_mlst_header()
#     full_header += get_rmp_mlst_header()
# 
#     full_header.append('spurious_virulence_hits')
# 
#     return stdout_header, full_header, res_headers
# 
# def get_resistance_results(data_folder, contigs, args, res_headers, kp_complex):
#     if not pathlib.Path(contigs).is_file():
#         raise OSError
#     if args.resistance:
#         gene_info, _, _ = read_class_file(data_folder + '/CARD_AMR_clustered.csv')
# 
#         # Only do mutation/truncation tests for Kp complex species.
#         if kp_complex:
#             qrdr = data_folder + '/QRDR_120.fasta'
#             trunc = data_folder + '/MgrB_and_PmrB.fasta'
#             omp = data_folder + '/OmpK.fasta'
#         else:
#             qrdr, trunc, omp = None, None, None
# 
#         seqs = data_folder + '/CARD_v3.0.8.fasta'
#         res_hits = resblast_one_assembly(contigs, gene_info, qrdr, trunc, omp, seqs,
#                                          args.min_coverage,  args.min_identity,
#                                          args.min_spurious_coverage, args.min_spurious_identity)
# 
#         # Double check that there weren't any results without a corresponding output header.
#         for h in res_hits.keys():
#             if h not in res_headers:
#                 sys.exit( f'Error: results contained a value ({h}) that is not covered by the '
#                           f'output headers')
# 
#         return {r: ';'.join(sorted(res_hits[r])) if r in res_hits else '-' for r in res_headers}
#     else:
#         return {}


import collections
import unittest
from kleborate.__main__ import get_output_headers, get_resistance_results
"""
Copyright 2020 Kat Holt
Copyright 2020 Ryan Wick (rrwick@gmail.com)
https://github.com/katholt/Kleborate/

This file is part of Kleborate. Kleborate is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by the Free Software Foundation,
either version 3 of the License, or (at your option) any later version. Kleborate is distributed in
the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
details. You should have received a copy of the GNU General Public License along with Kleborate. If
not, see <http://www.gnu.org/licenses/>.
"""




class TestResOmp(unittest.TestCase):
    """
    Tests calling of carbapenem resistance via the OmpK35/OmpK36 genes.
    """

    def setUp(self):
        self.data_dir = 'test/test_res_omp/data'
        Args = collections.namedtuple('Args', ['resistance', 'kaptive_k', 'kaptive_o',
                                               'min_coverage', 'min_identity',
                                               'min_spurious_coverage', 'min_spurious_identity'])
        self.args = Args(resistance=True, kaptive_k=False, kaptive_o=False,
                         min_coverage=80.0, min_identity=90.0,
                         min_spurious_coverage=40.0, min_spurious_identity=80.0)

# Path: kleborate/__main__.py
# def get_output_headers(args, data_folder):
#     """
#     There are two levels of output:
#       * stdout is simpler and displayed to the console
#       * full contains more and is saved to file
#     This function returns headers for both. It also returns the resistance headers in a separate
#     list, as they are used to total up some resistance summaries.
#     """
#     stdout_header = ['strain', 'species']
#     full_header = ['strain', 'species', 'species_match']
#     stdout_header += ['ST', 'virulence_score']
#     full_header += ['contig_count', 'N50', 'largest_contig', 'total_size', 'ambiguous_bases',
#                     'QC_warnings', 'ST', 'virulence_score']
# 
#     if args.resistance:
#         stdout_header.append('resistance_score')
#         full_header.append('resistance_score')
#         full_header.append('num_resistance_classes')
#         full_header.append('num_resistance_genes')
# 
#     other_columns = ['Yersiniabactin', 'YbST',
#                      'Colibactin', 'CbST',
#                      'Aerobactin', 'AbST',
#                      'Salmochelin', 'SmST',
#                      'RmpADC', 'RmST',
#                      'rmpA2',
#                      'wzi', 'K_locus']
#     stdout_header += other_columns
#     full_header += other_columns
# 
#     if args.kaptive_k:
#         stdout_header.append('K_locus_confidence')
#         full_header.append('K_type')
#         full_header.append('K_locus_problems')
#         full_header.append('K_locus_confidence')
#         full_header.append('K_locus_identity')
#         full_header.append('K_locus_missing_genes')
# 
#     if args.kaptive_o:
#         stdout_header.append('O_locus')
#         stdout_header.append('O_locus_confidence')
#         full_header.append('O_locus')
#         full_header.append('O_type')
#         full_header.append('O_locus_problems')
#         full_header.append('O_locus_confidence')
#         full_header.append('O_locus_identity')
#         full_header.append('O_locus_missing_genes')
# 
#     if args.resistance:
#         gene_info, res_classes, bla_classes = \
#             read_class_file(data_folder + '/CARD_AMR_clustered.csv')
#         res_headers = get_res_headers(res_classes, bla_classes)
#         res_headers += ['truncated_resistance_hits', 'spurious_resistance_hits']
#         stdout_header += res_headers
#         full_header += res_headers
#     else:
#         res_headers = []
# 
#     full_header.append('Chr_ST')
#     full_header += get_chromosome_mlst_header()
#     full_header += get_ybt_mlst_header()
#     full_header += get_clb_mlst_header()
#     full_header += get_iuc_mlst_header()
#     full_header += get_iro_mlst_header()
#     full_header += get_rmp_mlst_header()
# 
#     full_header.append('spurious_virulence_hits')
# 
#     return stdout_header, full_header, res_headers
# 
# def get_resistance_results(data_folder, contigs, args, res_headers, kp_complex):
#     if not pathlib.Path(contigs).is_file():
#         raise OSError
#     if args.resistance:
#         gene_info, _, _ = read_class_file(data_folder + '/CARD_AMR_clustered.csv')
# 
#         # Only do mutation/truncation tests for Kp complex species.
#         if kp_complex:
#             qrdr = data_folder + '/QRDR_120.fasta'
#             trunc = data_folder + '/MgrB_and_PmrB.fasta'
#             omp = data_folder + '/OmpK.fasta'
#         else:
#             qrdr, trunc, omp = None, None, None
# 
#         seqs = data_folder + '/CARD_v3.0.8.fasta'
#         res_hits = resblast_one_assembly(contigs, gene_info, qrdr, trunc, omp, seqs,
#                                          args.min_coverage,  args.min_identity,
#                                          args.min_spurious_coverage, args.min_spurious_identity)
# 
#         # Double check that there weren't any results without a corresponding output header.
#         for h in res_hits.keys():
#             if h not in res_headers:
#                 sys.exit( f'Error: results contained a value ({h}) that is not covered by the '
#                           f'output headers')
# 
#         return {r: ';'.join(sorted(res_hits[r])) if r in res_hits else '-' for r in res_headers}
#     else:
#         return {}


# Path: test/test_res_omp.py
import collections
import unittest
from kleborate.__main__ import get_output_headers, get_resistance_results

"""
Copyright 2020 Kat Holt
Copyright 2020 Ryan Wick (rrwick@gmail.com)
https://github.com/katholt/Kleborate/

This file is part of Kleborate. Kleborate is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by the Free Software Foundation,
either version 3 of the License, or (at your option) any later version. Kleborate is distributed in
the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
details. You should have received a copy of the GNU General Public License along with Kleborate. If
not, see <http://www.gnu.org/licenses/>.
"""




class TestResOmp(unittest.TestCase):
    """
    Tests calling of carbapenem resistance via the OmpK35/OmpK36 genes.
    """

    def setUp(self):
        self.data_dir = 'test/test_res_omp/data'
        Args = collections.namedtuple('Args', ['resistance', 'kaptive_k', 'kaptive_o',
                                               'min_coverage', 'min_identity',
                                               'min_spurious_coverage', 'min_spurious_identity'])
        self.args = Args(resistance=True, kaptive_k=False, kaptive_o=False,
                         min_coverage=80.0, min_identity=90.0,
                         min_spurious_coverage=40.0, min_spurious_identity=80.0)

        _, _, self.res_headers = get_output_headers(self.args, self.data_dir)