import os
import unittest


class FastaTester(unittest.TestCase):
    def setUp(self) -> None:
        from treesapp.fasta import load_fasta_header_regexes
        from .testing_utils import get_test_data
        self.header_regexes = load_fasta_header_regexes(code_name="PuhA")

        self.test_fa = get_test_data("McrA_eval.faa")
        self.test_fq = get_test_data("test_TarA.1.fq")
        self.troublesome_fa = get_test_data("fasta_parser_test.fasta")
        return

    def tearDown(self) -> None:
        from glob import glob
        test_fa_prefix, _ = os.path.splitext(os.path.basename(self.test_fa))
        test_fq_prefix, _ = os.path.splitext(os.path.basename(self.test_fq))
        for file_name in sorted(glob(test_fa_prefix + "*") + glob(test_fq_prefix + "*")):
            if os.path.isfile(file_name):
                os.remove(file_name)
        return

    def test_get_header_format(self):
        from treesapp.fasta import get_header_format
        _, header_db, _ = get_header_format("650716056_50856936 Metig_1237", self.header_regexes)
        self.assertEqual("unformatted", header_db)
        seq_name = ">1000565.METUNv1_03969"
        _, header_db, _ = get_header_format(seq_name, self.header_regexes)
        self.assertEqual("eggnog", header_db)
        return

    def test_sequence_info_groups(self):
        from treesapp.fasta import get_header_format
        from treesapp.fasta import sequence_info_groups

        # Test one: unformatted
        seq_name = "650716056_50856936 Metig_1237"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual("650716056_50856936", seq_info_tuple.accession)

        # Test two: NCBI protein
        seq_name = ">AKB49151.1 Methyl coenzyme M reductase alpha subunit [Methanosarcina sp. Kolksee]"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual("AKB49151", seq_info_tuple.accession)
        seq_name = ">WP_048080940.1"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual("WP_048080940", seq_info_tuple.accession)
        self.assertEqual("WP_048080940.1", seq_info_tuple.version)

        # Test three: EggNOG
        seq_name = ">1000565.METUNv1_03969"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual("1000565.METUNv1_03969", seq_info_tuple.accession)
        self.assertEqual("1000565", seq_info_tuple.taxid)
        seq_name = "1148.1006613"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual(seq_name, seq_info_tuple.version)
        self.assertEqual("1148", seq_info_tuple.taxid)
        seq_name = "4155.Migut.B00751.1.p"
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        self.assertEqual("eggnog", header_db)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual(seq_name, seq_info_tuple.version)
        self.assertEqual("4155", seq_info_tuple.taxid)

        # Test four: SwissProt, representing accessions with database prefix (e.g. 'sp') with accession enclosed by '|'
        seq_name = 'sp|P06008|RCEH_BLAVI Reaction center protein H chain OS=Blastochloris viridis GN=puhA'
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        self.assertEqual("sp", header_db)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name)
        self.assertEqual("P06008", seq_info_tuple.accession)
        return

    def test_sequence_info_groups_assigned(self):
        from treesapp.fasta import get_header_format, sequence_info_groups
        seq_name = 'sp|P06008|RCEH_BLAVI Reaction center protein H chain OS=Blastochloris viridis|PuhA|1_258'
        header_format_re, header_db, _ = get_header_format(seq_name, self.header_regexes)
        self.assertEqual("ts_assign", header_db)
        seq_info_tuple = sequence_info_groups(header_format_re.match(seq_name), header_db, seq_name, self.header_regexes)
        self.assertEqual("P06008", seq_info_tuple.accession)
        self.assertEqual("sp|P06008|RCEH_BLAVI", seq_info_tuple.version)
        self.assertEqual("sp|P06008|RCEH_BLAVI Reaction center protein H chain OS=Blastochloris viridis|PuhA|1_258",
                         seq_info_tuple.description)
        return

    def test_read_fasta_to_dict(self):
        from treesapp.fasta import read_fasta_to_dict

        # Test normal functionality with a fasta file
        fa_dict = read_fasta_to_dict(self.test_fa)
        self.assertEqual(236, len(fa_dict))

        # Test reading a file that doesn't exist
        self.assertEqual({}, read_fasta_to_dict("./fake_file.fa"))
        return

    def test_split_fa(self):
        from treesapp.fasta import split_fa, read_fasta_to_dict
        # Test splitting by number of files
        split_files = split_fa(fastx=self.test_fa, outdir="./", file_num=4)
        num_split_seqs = sum([len(read_fasta_to_dict(f)) for f in split_files])
        self.assertEqual(len(read_fasta_to_dict(self.test_fa)), num_split_seqs)
        self.assertEqual(4, len(split_files))

        # Test splitting by count
        split_files = split_fa(fastx=self.test_fa, outdir="./", file_num=1, max_seq_count=236)
        self.assertEqual(1, len(split_files))
        split_files = split_fa(fastx=self.test_fa, outdir="./", file_num=1, max_seq_count=24)
        num_split_seqs = sum([len(read_fasta_to_dict(f)) for f in split_files])
        self.assertEqual(len(read_fasta_to_dict(self.test_fa)), num_split_seqs)
        self.assertEqual(10, len(split_files))
        return

    def test_parameterize_sub_file_size(self):
        from treesapp.fasta import parameterize_sub_file_size
        fsize, nseqs = parameterize_sub_file_size(file_name=self.test_fq, num_files=4, max_seqs=0, record_len=4)
        self.assertEqual(1032, fsize)
        self.assertEqual(0, nseqs)
        fsize, nseqs = parameterize_sub_file_size(file_name=self.test_fa, num_files=0, max_seqs=10)
        self.assertEqual(10, nseqs)
        self.assertEqual(0, fsize)
        return

    def test_fq2fa(self):
        from treesapp.fasta import fq2fa
        split_files = fq2fa(fastx=self.test_fq, outdir="./", file_num=4)
        line_num = 0
        for f in split_files:
            with open(f, 'r') as f_handler:
                line_num += len(f_handler.readlines())
        self.assertEqual(24, line_num)
        self.assertEqual(4, len(split_files))
        split_files = fq2fa(fastx=self.test_fq, outdir="./", max_seq_count=6)
        for f_path in split_files:
            self.assertTrue(os.path.isfile(f_path))
        return

    def test_format_read_fasta(self):
        from treesapp.fasta import format_read_fasta
        fasta_dict = format_read_fasta(fasta_input=self.troublesome_fa, molecule="prot", min_seq_length=50)
        self.assertEqual(5, len(fasta_dict))
        return

    def test_trim_to_length(self):
        from treesapp.fasta import FASTA, register_headers
        _test_length = 10
        mock_fa = FASTA("mock.fa")
        mock_fa.fasta_dict = {"seq_1": "ADJKHJHFAJKHDFKJH", "seq_2": "MCNMVBASUIQRUIYA", "seq_3": "AAAAAAA"}
        mock_fa.header_registry = register_headers(list(mock_fa.fasta_dict.keys()))

        mock_fa.trim_to_length(_test_length)
        seq_lengths = [len(seq) for seq in mock_fa.fasta_dict.values()]
        # Check sequence lengths are correct
        self.assertTrue(min(seq_lengths) == max(seq_lengths) == _test_length)
        # Check FASTA attributes
        self.assertEqual(["seq_1", "seq_2"], mock_fa.get_seq_names())
        self.assertEqual(2, mock_fa.n_seqs())
        self.assertEqual(2, len(mock_fa.header_registry))
        return

    def test_remove_shorter_than(self):
        from treesapp.fasta import FASTA
        test_fa = FASTA(self.test_fa)
        test_fa.load_fasta()
        self.assertEqual(0, test_fa.remove_shorter_than(min_len=30))
        self.assertEqual(1, test_fa.remove_shorter_than(min_len=1000))
        return

    def test_clear(self):
        from treesapp.fasta import FASTA
        test_fa = FASTA(self.test_fa)
        test_fa.load_fasta()
        test_fa.clear()
        self.assertEqual({}, test_fa.fasta_dict)
        return


class HeaderTester(unittest.TestCase):
    def setUp(self) -> None:
        self.eggnog_one = '1148.1006613'
        self.eggnog_two = '28072.Nos7524_3177'
        self.eggnog_three = '59919.PMM0223'

    def test_find_accession(self):
        from treesapp.fasta import Header
        h = Header(self.eggnog_one)
        h.find_accession()
        self.assertEqual('1148.1006613', h.accession)
        self.assertEqual('1148.1006613', h.version)

        h.original = self.eggnog_two
        h.find_accession()
        self.assertEqual(self.eggnog_two, h.version)

        h.original = self.eggnog_three
        h.find_accession()
        self.assertEqual(self.eggnog_three, h.version)
        return

    def test_get_info(self):
        from treesapp.fasta import Header
        h = Header(self.eggnog_one)
        info = h.get_info()
        self.assertEqual(3, len(info.split('\n')))
        return


if __name__ == '__main__':
    unittest.main()
