import unittest
import urllib.request
import os
import zipfile

from Bio.SeqRecord import SeqRecord

class AlignmentSeqTestCase(unittest.TestCase):
    def setUp(self):
        balibase_zippath = "balibase.zip"
        balibase_path = "balibase"
        testfile_path = os.path.join(
            balibase_path, 
            "RV11.unaligned", "BBS11001.fasta")

        if not os.path.isdir(balibase_path):
            os.mkdir(balibase_path)
        if not os.path.isfile(testfile_path):
            if not os.path.isfile(balibase_zippath):
                print("Fetching balibase archive from moodle...")
                urllib.request.urlretrieve(
                    "https://moodle.polytechnique.fr/mod/resource/view.php?id=38570",
                    balibase_zippath)
            with zipfile.ZipFile(balibase_zippath) as balibase_zip:
                balibase_zip.extractall()

    def test_simple_align(self):
        from alignementseq import align
        score_fn = lambda a,b : -2 if a == '' or b == '' else -1 if a != b else 1

        cases = [
                ("CHAT", "CAT", 1, "CHAT", "C-AT"),
                ("CHAT", "CGAT", 2, "CHAT", "CGAT"),
                ("CHAT", "AT", -2, "CHAT", "--AT")
                ]
        for s1, s2, exp_score, exp_r1, exp_r2 in cases:
            score, r1, r2 = align(SeqRecord(s1), SeqRecord(s2), score_fn)
            self.assertEqual(score, exp_score)
            self.assertEqual(r1.seq, exp_r1)
            self.assertEqual(r2.seq, exp_r2)

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