# Copyright 2001 by Brad Chapman.  All rights reserved.
# Revisions copyright 2008-2009 by Michiel de Hoon. All rights reserved.
# Revisions copyright 2010-2011 by Peter Cock. All rights reserved.
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

"""Tests for Primer-based programs in the Emboss suite."""

# standard library
import os
import unittest

# local stuff
from Bio.Emboss import PrimerSearch, Primer3


class Primer3ParseTest(unittest.TestCase):
    def setUp(self):
        self.test_files = [
            os.path.join("Emboss", "bac_find.primer3"),
            os.path.join("Emboss", "cds_forward.primer3"),
            os.path.join("Emboss", "cds_reverse.primer3"),
            os.path.join("Emboss", "short.primer3"),
            os.path.join("Emboss", "internal_oligo.primer3"),
            os.path.join("Emboss", "no_oligo.primer3"),
        ]

    def test_simple_parse(self):
        """Make sure that we can use all single target primer3 files."""
        for file in self.test_files:
            # First using read...
            with open(file) as handle:
                Primer3.read(handle)
            # Now using parse...
            with open(file) as handle:
                self.assertEqual(1, len(list(Primer3.parse(handle))))

    def test_indepth_regular_parse(self):
        """Make sure we get the data from normal primer3 files okay."""
        regular_file = self.test_files[0]
        with open(regular_file) as handle:
            primer_info = Primer3.read(handle)

        self.assertEqual(len(primer_info.primers), 5)
        self.assertEqual(primer_info.comments, "# PRIMER3 RESULTS FOR AC074298\n")
        self.assertEqual(primer_info.primers[1].forward_seq, "CCGGTTTCTCTGGTTGAAAA")
        self.assertEqual(primer_info.primers[2].reverse_seq, "TCACATTCCCAAATGTAGATCG")
        self.assertEqual(primer_info.primers[0].size, 218)
        self.assertEqual(len(primer_info.primers[0]), 218)
        self.assertEqual(primer_info.primers[3].forward_start, 112)
        self.assertEqual(primer_info.primers[3].forward_length, 20)
        self.assertEqual(primer_info.primers[3].forward_tm, 59.57)
        self.assertEqual(primer_info.primers[3].forward_gc, 45.00)
        self.assertEqual(primer_info.primers[4].reverse_start, 304)
        self.assertEqual(primer_info.primers[4].reverse_length, 22)
        self.assertEqual(primer_info.primers[4].reverse_tm, 59.61)
        self.assertEqual(primer_info.primers[4].reverse_gc, 40.91)

    def test_in_depth_single_parse(self):
        """Make sure we get info right from a single primer find."""
        file = self.test_files[1]
        with open(file) as handle:
            primer_info = Primer3.read(handle)

        self.assertEqual(len(primer_info.primers), 5)
        self.assertEqual(primer_info.comments, "# PRIMER3 RESULTS FOR 26964-28647#\n")
        self.assertEqual(primer_info.primers[1].reverse_seq, "")
        self.assertEqual(primer_info.primers[1].internal_seq, "")
        self.assertEqual(primer_info.primers[3].forward_seq, "TGTGATTGCTTGAGCTGGAC")
        self.assertEqual(primer_info.primers[3].internal_seq, "")
        self.assertEqual(primer_info.primers[3].forward_start, 253)

    def test_internal_oligo_single_parse(self):
        """Make sure we can parse an internal oligo file correctly."""
        # these files are generated when designing hybridization probes.
        file = self.test_files[4]
        with open(file) as handle:
            primer_info = Primer3.read(handle)

        self.assertEqual(len(primer_info.primers), 5)
        self.assertEqual(primer_info.comments, "# EPRIMER3 RESULTS FOR YNL138W-A\n")
        self.assertEqual(primer_info.primers[0].internal_length, 22)
        self.assertEqual(primer_info.primers[1].internal_seq, "TTGCGCTTTAGTTTGAATTGAA")
        self.assertEqual(primer_info.primers[2].internal_tm, 58.62)
        self.assertEqual(primer_info.primers[3].internal_start, 16)
        self.assertEqual(primer_info.primers[4].internal_gc, 35.00)

    def test_multi_record_fwd(self):
        """Test parsing multiple primer sets (NirK forward)."""
        with open(os.path.join("Emboss", "NirK.primer3")) as handle:
            targets = list(Primer3.parse(handle))

        self.assertEqual(len(targets), 16)
        for target in targets:
            self.assertEqual(len(target.primers), 5)

        self.assertEqual(targets[0].primers[0].forward_seq, "GCAAACTGAAAAGCGGACTC")
        self.assertEqual(targets[0].primers[1].forward_seq, "GGGACGTACTTTCGCACAAT")
        self.assertEqual(targets[0].primers[2].forward_seq, "GTCTTATGCGTGGTGGAGGT")
        self.assertEqual(targets[0].primers[3].forward_seq, "GTACATCAACATCCGCAACG")
        self.assertEqual(targets[0].primers[4].forward_seq, "CGTACATCAACATCCGCAAC")

        self.assertEqual(targets[1].primers[0].forward_seq, "GGAAGTGCTTCTCGTTTTCG")
        self.assertEqual(targets[1].primers[1].forward_seq, "TACAGAGCGTCACGGATGAG")
        self.assertEqual(targets[1].primers[2].forward_seq, "TTGTCATCGTGCTCTTCGTC")
        self.assertEqual(targets[1].primers[3].forward_seq, "GACTCCAACCTCAGCTTTCG")
        self.assertEqual(targets[1].primers[4].forward_seq, "GGCACGAAGAAGGACAGAAG")

        self.assertEqual(targets[15].primers[0].forward_seq, "TGCTTGAAAATGACGCACTC")
        self.assertEqual(targets[15].primers[1].forward_seq, "CTCGCTGGCTAGGTCATAGG")
        self.assertEqual(targets[15].primers[2].forward_seq, "TATCGCACCAAACACGGTAA")
        self.assertEqual(targets[15].primers[3].forward_seq, "CGATTACCCTCACCGTCACT")
        self.assertEqual(targets[15].primers[4].forward_seq, "TATCGCAACCACTGAGCAAG")

    def test_multi_record_full(self):
        """Test parsing multiple primer sets (NirK full)."""
        with open(os.path.join("Emboss", "NirK_full.primer3")) as handle:
            targets = list(Primer3.parse(handle))

        self.assertEqual(len(targets), 16)
        for target in targets:
            self.assertEqual(len(target.primers), 5)

        self.assertEqual(targets[15].primers[0].forward_seq, "ACTCACTTCGGCTGAATGCT")
        self.assertEqual(targets[15].primers[1].forward_seq, "GGCGATTAGCGCTGTCTATC")
        self.assertEqual(targets[15].primers[2].forward_seq, "ACTCACTTCGGCTGAATGCT")
        self.assertEqual(targets[15].primers[3].forward_seq, "TAGGCGTATAGACCGGGTTG")
        self.assertEqual(targets[15].primers[4].forward_seq, "AGCAAGCTGACCACTGGTTT")

        self.assertEqual(targets[15].primers[0].reverse_seq, "CATTTAATCCGGATGCCAAC")
        self.assertEqual(targets[15].primers[1].reverse_seq, "TGGCCTTTCTCTCCTCTTCA")
        self.assertEqual(targets[15].primers[2].reverse_seq, "ATTTAATCCGGATGCCAACA")
        self.assertEqual(targets[15].primers[3].reverse_seq, "CACACATTATTGGCGGTCAC")
        self.assertEqual(targets[15].primers[4].reverse_seq, "TCTGAAACCACCAAGGAAGC")

        self.assertEqual(targets[15].primers[0].internal_seq, "CCCACCAATATTTGGCTAGC")
        self.assertEqual(targets[15].primers[1].internal_seq, "AATCTTCTGTGCACCTTGCC")
        self.assertEqual(targets[15].primers[2].internal_seq, "CCCACCAATATTTGGCTAGC")
        self.assertEqual(targets[15].primers[3].internal_seq, "TGAGCCTGTGTTCCACACAT")
        self.assertEqual(targets[15].primers[4].internal_seq, "CTATGCCCTTCTGCCACAAT")


class PrimersearchParseTest(unittest.TestCase):
    def setUp(self):
        self.test_files = [os.path.join("Emboss", "bac_find.psearch")]

    def test_simple_parse(self):
        """Make sure that we can parse all primersearch files."""
        for file in self.test_files:
            with open(file) as handle:
                PrimerSearch.read(handle)

    def test_in_depth_normal_parse(self):
        """Make sure the output from a simple primersearch file is correct."""
        file = self.test_files[0]
        with open(file) as handle:
            amp_info = PrimerSearch.read(handle)

        self.assertEqual(len(amp_info.amplifiers), 1)
        self.assertIn("Test", amp_info.amplifiers)
        self.assertEqual(len(amp_info.amplifiers["Test"]), 1)

        self.assertEqual(amp_info.amplifiers["Test"][0].length, 218)
        self.assertEqual(
            amp_info.amplifiers["Test"][0].hit_info,
            "AC074298 AC074298 \n"
            "\tTelomere associated sequence for Arabidopsis thaliana "
            "TEL1N from chromosome I, complete sequence.\n"
            "\tCCGGTTTCTCTGGTTGAAAA hits forward strand at 114 "
            "with 0 mismatches\n"
            "\tTCACATTCCCAAATGTAGATCG hits reverse strand at "
            "[114] with 0 mismatches",
        )


class PrimerSearchInputTest(unittest.TestCase):
    """Test creating input files for primersearch."""

    def setUp(self):
        pass

    def test_primer_representation(self):
        """Make sure we can output primer information correctly."""
        p_info = PrimerSearch.InputRecord()
        p_info.add_primer_set("Test", "GATC", "CATG")
        p_info.add_primer_set("Test2", "AATA", "TTAT")

        output = str(p_info)
        self.assertEqual(output, "Test GATC CATG\nTest2 AATA TTAT\n")


if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)
