#!/usr/bin/env python

# 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.

"""Test the different representations of Genes.

This exercises the Motif, Schema and Signature methods of representing
genes, as well as generic Pattern methods.
"""
# standard library
from __future__ import print_function

import os
import unittest

# Biopython
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.Alphabet import IUPAC

# stuff we are testing
from Bio.NeuralNetwork.Gene import Schema
from Bio.NeuralNetwork.Gene import Motif
from Bio.NeuralNetwork.Gene import Signature
from Bio.NeuralNetwork.Gene import Pattern

VERBOSE = 0


# --- Tests for Pattern

class PatternIOTest(unittest.TestCase):
    """Tests for reading and writing patterns to a file.
    """
    def setUp(self):
        self.alphabet = IUPAC.ambiguous_dna
        self.test_file = os.path.join("NeuralNetwork", "patternio.txt")
        # Remove any existing copy of the output file,
        if os.path.isfile(self.test_file):
            os.remove(self.test_file)
        self.pattern_io = Pattern.PatternIO(self.alphabet)

    def tearDown(self):
        # Clean up by removing our output file,
        if os.path.isfile(self.test_file):
            os.remove(self.test_file)

    def test_motif(self):
        """Reading and writing motifs to a file
        """
        # write to a file
        motifs = ["GAC", "AAA", "TTT", "GGG"]
        with open(self.test_file, "w") as output_handle:
            self.pattern_io.write(motifs, output_handle)

        # read 'em back
        with open(self.test_file, "r") as input_handle:
            read_motifs = self.pattern_io.read(input_handle)
        self.assertEqual(read_motifs, motifs,
                         "Failed to get back expected motifs %s, got %s"
                         % (motifs, read_motifs))

        # write seqs
        seq_motifs = []
        for motif in motifs:
            seq_motifs.append(Seq(motif, self.alphabet))
        with open(self.test_file, "w") as output_handle:
            self.pattern_io.write_seq(seq_motifs, output_handle)

        # read the seqs back
        with open(self.test_file, "r") as input_handle:
            read_motifs = self.pattern_io.read(input_handle)
        self.assertEqual(read_motifs, motifs,
                         "Failed to get back expected motifs %s from seqs, got %s"
                         % (motifs, read_motifs))

    def test_schema(self):
        """Reading and writing schemas to a file.
        """
        schemas = ["GTR", "GAC"]
        # write out the schemas
        with open(self.test_file, "w") as output_handle:
            self.pattern_io.write(schemas, output_handle)

        # read back the schemas
        with open(self.test_file, "r") as input_handle:
            read_schemas = self.pattern_io.read(input_handle)
        self.assertEqual(schemas, read_schemas,
                             "Read incorrect schemas %s, expected %s."
                             % (read_schemas, schemas))

        # --- make sure inappropriate alphabets are reported
        schemas = ["GTR", "G*C"]  # '*' not in the unambigous alphabet
        with open(self.test_file, "w") as output_handle:
            self.pattern_io.write(schemas, output_handle)

        input_handle = open(self.test_file, "r")
        try:
            read_schemas = self.pattern_io.read(input_handle)
            raise AssertionError("Did not report error on bad alphabet.")
        except ValueError:
            pass  # expected behavior
        except:
            raise AssertionError("Got unexpected error while reading.")
        input_handle.close()

    def test_signature(self):
        """Reading and writing signatures to a file.
        """
        signatures = [("GAC", "GAC"), ("AAA", "TTT")]
        with open(self.test_file, "w") as output_handle:
            self.pattern_io.write(signatures, output_handle)

        input_handle = open(self.test_file, "r")
        read_sigs = self.pattern_io.read(input_handle)
        input_handle.close()
        self.assertEqual(read_sigs, signatures,
                         "Got back unexpected signatures %s, wanted %s"
                         % (read_sigs, signatures))


class PatternRepositoryTest(unittest.TestCase):
    """Tests for retrieving info from a repository of patterns.
    """
    def setUp(self):
        self.motifs = {"GATC": 30,
                       "GGGG": 10,
                       "GTAG": 0,
                       "AAAA": -10,
                       "ATAT": -20}

        self.repository = Pattern.PatternRepository(self.motifs)

    def test_get_all(self):
        """Retrieve all patterns from a repository.
        """
        all_motifs = self.repository.get_all()

        self.assertEqual(all_motifs,
                         ["GATC", "GGGG", "GTAG", "AAAA", "ATAT"],
                         "Unexpected motifs returned %s" % all_motifs)

    def test_get_random(self):
        """Retrieve random patterns from the repository.
        """
        for num_patterns in range(5):
            patterns = self.repository.get_random(num_patterns)
            self.assertEqual(len(patterns), num_patterns,
                             "Got unexpected number of patterns %s, expected %s"
                             % (len(patterns), num_patterns))

            for pattern in patterns:
                self.assertTrue(pattern in list(self.motifs.keys()),
                                "Got unexpected pattern %s" % pattern)

    def test_get_top_percentage(self):
        """Retrieve the top percentge of patterns from the repository.
        """
        for num_patterns, percentage in ((1, 0.2), (2, .4), (5, 1.0)):
            patterns = self.repository.get_top_percentage(percentage)
            self.assertEqual(len(patterns), num_patterns,
                             "Got unexpected number of patterns %s, expected %s"
                             % (len(patterns), num_patterns))

            for pattern in patterns:
                self.assertTrue(pattern in list(self.motifs.keys()),
                                "Got unexpected pattern %s" % pattern)

    def test_get_top(self):
        """Retrieve a certain number of the top patterns.
        """
        for num_patterns in range(5):
            patterns = self.repository.get_top(num_patterns)
            self.assertEqual(len(patterns), num_patterns,
                             "Got unexpected number of patterns %s, expected %s"
                             % (len(patterns), num_patterns))

            for pattern in patterns:
                self.assertTrue(pattern in list(self.motifs.keys()),
                                "Got unexpected pattern %s" % pattern)

    def test_get_differing(self):
        """Retrieve patterns from both sides of the list (top and bottom).
        """
        patterns = self.repository.get_differing(2, 2)
        self.assertEqual(patterns,
                         ["GATC", "GGGG", "AAAA", "ATAT"],
                         "Got unexpected patterns %s" % patterns)

    def test_remove_polyA(self):
        """Test the ability to remove A rich patterns from the repository.
        """
        patterns = self.repository.get_all()
        self.assertEqual(len(patterns), 5,
                         "Unexpected starting: %s" % patterns)

        self.repository.remove_polyA()

        patterns = self.repository.get_all()
        self.assertEqual(len(patterns), 3,
                         "Unexpected ending: %s" % patterns)
        self.assertEqual(patterns,
                         ["GATC", "GGGG", "GTAG"],
                         "Unexpected patterns: %s" % patterns)

    def test_count(self):
        """Retrieve counts for particular patterns in the repository.
        """
        num_times = self.repository.count("GGGG")
        self.assertEqual(num_times, 10,
                         "Did not count item in the respository: %s" % num_times)

        num_times = self.repository.count("NOT_IN_THERE")
        self.assertEqual(num_times, 0,
                         "Counted items not in repository: %s" % num_times)


# --- Tests for motifs

class MotifFinderTest(unittest.TestCase):
    """Tests for finding motifs from sequences.
    """
    def setUp(self):
        test_file = os.path.join('NeuralNetwork', 'enolase.fasta')
        diff_file = os.path.join('NeuralNetwork', 'repeat.fasta')

        self.test_records = []
        self.diff_records = []

        # load the records
        for file, records in ((test_file, self.test_records),
                              (diff_file, self.diff_records)):

            handle = open(file, 'r')

            iterator = SeqIO.parse(handle, "fasta",
                                   alphabet=IUPAC.unambiguous_dna)
            while True:
                try:
                    seq_record = next(iterator)
                except StopIteration:
                    break
                if seq_record is None:
                    break

                records.append(seq_record)

            handle.close()

        self.motif_finder = Motif.MotifFinder()

    def test_find(self):
        """Find all motifs in a set of sequences.
        """
        motif_repository = self.motif_finder.find(self.test_records, 8)
        top_motif = motif_repository.get_top(1)

        self.assertEqual(top_motif[0], 'TTGGAAAG',
                         "Got unexpected motif %s" % top_motif[0])

    def test_find_differences(self):
        """Find the difference in motif counts between two sets of sequences.
        """
        motif_repository = \
               self.motif_finder.find_differences(self.test_records,
                                                  self.diff_records, 8)

        top, bottom = motif_repository.get_differing(1, 1)

        self.assertEqual(top, "TTGGAAAG",
                         "Got unexpected top motif %s" % top)
        self.assertEqual(bottom, "AATGGCAT",
                         "Got unexpected bottom motif %s" % bottom)


class MotifCoderTest(unittest.TestCase):
    """Test the ability to encode sequences as a set of motifs.
    """
    def setUp(self):
        motifs = ["GAG", "GAT", "GCC", "ATA"]

        self.match_strings = (("GATCGCC", [0.0, 1.0, 1.0, 0.0]),
                              ("GATGATCGAGCC", [.5, 1.0, .5, 0.0]))

        self.coder = Motif.MotifCoder(motifs)

    def test_representation(self):
        """Convert a sequence into its motif representation.
        """
        for match_string, expected in self.match_strings:
            seq_to_code = Seq(match_string, IUPAC.unambiguous_dna)
            matches = self.coder.representation(seq_to_code)

            self.assertEqual(matches, expected,
                             "Did not match representation, expected %s, got %s"
                             % (expected, matches))


# --- Tests for schemas

class SchemaTest(unittest.TestCase):
    """Matching ambiguous motifs with multiple ambiguity characters.
    """
    def setUp(self):
        ambiguity_chars = {"G": "G",
                           "A": "A",
                           "T": "T",
                           "C": "C",
                           "R": "AG",
                           "*": "AGTC"}

        self.motif_coder = Schema.Schema(ambiguity_chars)

        self.match_string = "GATAG"
        self.match_info = [("GA", ["GA"]),
                           ("GATAG", ["GATAG"]),
                           ("GA*AG", ["GATAG"]),
                           ("GATRG", ["GATAG"]),
                           ("*A", ["GA", "TA"])]

    def test_find_matches(self):
        """Find all matches in a sequence.
        """
        for motif, expected in self.match_info:
            found_matches = self.motif_coder.find_matches(motif,
                                                          self.match_string)
            self.assertEqual(found_matches, expected,
                             "Expected %s, got %s"
                             % (expected, found_matches))

    def test_num_matches(self):
        """Find how many matches are present in a sequence.
        """
        for motif, expected in self.match_info:
            num_matches = self.motif_coder.num_matches(motif,
                                                       self.match_string)
            self.assertEqual(num_matches, len(expected),
                             "Expected %s, got %s"
                             % (num_matches, len(expected)))

    def test_find_ambiguous(self):
        """Find the positions of ambiguous items in a sequence.
        """
        ambig_info = (("GATC", []),
                      ("G***", [1, 2, 3]),
                      ("GART", [2]),
                      ("*R*R", [0, 1, 2, 3]))

        for motif, expected in ambig_info:
            found_positions = self.motif_coder.find_ambiguous(motif)
            self.assertEqual(found_positions, expected,
                             "Expected %s, got %s for %s"
                             % (expected, found_positions, motif))

    def test_num_ambiguous(self):
        """Find the number of ambiguous items in a sequence.
        """
        ambig_info = (("GATC", 0),
                      ("G***", 3),
                      ("GART", 1),
                      ("*R*R", 4))

        for motif, expected in ambig_info:
            found_num = self.motif_coder.num_ambiguous(motif)
            self.assertEqual(found_num, expected,
                             "Expected %s, got %s for %s"
                             % (expected, found_num, motif))

    def test_motif_cache(self):
        """Make sure motif compiled regular expressions are cached properly.
        """
        test_motif = "GATC"

        self.motif_coder.find_matches(test_motif, "GATCGATC")

        self.assertTrue(test_motif in self.motif_coder._motif_cache,
                        "Did not find motif cached properly.")

        # make sure we don't bomb out if we use the same motif twice
        self.motif_coder.find_matches(test_motif, "GATCGATC")

    def test_all_unambiguous(self):
        """Return all unambiguous characters that can be in a motif.
        """
        found_unambig = self.motif_coder.all_unambiguous()

        expected = ["A", "C", "G", "T"]
        self.assertEqual(found_unambig, expected,
                         "Got %s, expected %s" % (found_unambig, expected))


class SchemaFinderTest(unittest.TestCase):
    """Test finding schemas from a set of sequences.
    """
    def setUp(self):
        test_file = os.path.join('NeuralNetwork', 'enolase.fasta')
        diff_file = os.path.join('NeuralNetwork', 'repeat.fasta')

        self.test_records = []
        self.diff_records = []

        # load the records
        for file, records in ((test_file, self.test_records),
                              (diff_file, self.diff_records)):

            handle = open(file, 'r')
            records.extend(SeqIO.parse(handle, "fasta",
                                       alphabet=IUPAC.unambiguous_dna))
            handle.close()

        self.num_schemas = 2
        schema_ga = Schema.GeneticAlgorithmFinder()
        schema_ga.min_generations = 1
        self.finder = Schema.SchemaFinder(num_schemas=self.num_schemas,
                                          schema_finder=schema_ga)

    def test_find(self):
        """Find schemas from sequence inputs.
        """
        # this test takes too long
        if VERBOSE:
            repository = self.finder.find(self.test_records + self.diff_records)
            schemas = repository.get_all()

            self.assertTrue(len(schemas) >= self.num_schemas,
                            "Got too few schemas.")

    def test_find_differences(self):
        """Find schemas that differentiate between two sets of sequences.
        """
        # this test takes too long
        if VERBOSE:
            repository = self.finder.find_differences(self.test_records,
                                                      self.diff_records)
            schemas = repository.get_all()

            self.assertTrue(len(schemas) >= self.num_schemas,
                            "Got too few schemas.")


class SchemaCoderTest(unittest.TestCase):
    """Test encoding sequences as a grouping of motifs.
    """
    def setUp(self):
        ambiguity_chars = {"G": "G",
                           "A": "A",
                           "T": "T",
                           "C": "C",
                           "R": "AG",
                           "*": "AGTC"}

        motif_representation = Schema.Schema(ambiguity_chars)
        motifs = ("GA", "GATAG", "GA*AG", "GATRG", "*A")
        self.motif_coder = Schema.SchemaCoder(motifs,
                                              motif_representation)

        self.match_strings = [("GATAG", [.5, .5, .5, .5, 1.0]),
                              ("GAGAGATA", [float(3) / float(4), 0,
                                            float(1) / float(4), 0,
                                            1])]

    def test_representation(self):
        """Convert a string into a representation of motifs.
        """
        for match_string, expected in self.match_strings:
            match_seq = Seq(match_string, IUPAC.unambiguous_dna)
            found_rep = self.motif_coder.representation(match_seq)
            self.assertEqual(found_rep, expected,
                             "Got %s, expected %s" % (found_rep, expected))


class SchemaMatchingTest(unittest.TestCase):
    """Matching schema to strings works correctly."""

    def shortDescription(self):
        return "%s:%s" % (self.__class__.__name__, self.__doc__)

    def runTest(self):
        match = Schema.matches_schema("GATC", "AAAAA")
        self.assertEqual(match, 0,
                         "Expected no match because of length differences")

        match = Schema.matches_schema("GATC", "GAT*")
        self.assertEqual(match, 1,
                         "Expected match")

        match = Schema.matches_schema("GATC", "GATC")
        self.assertEqual(match, 1,
                         "Expected match")

        match = Schema.matches_schema("GATC", "C*TC")
        self.assertEqual(match, 0,
                         "Expected no match because of char mismatch.")

        match = Schema.matches_schema("G*TC", "*TTC")
        self.assertEqual(match, 1,
                         "Expected match because of ambiguity.")


class SchemaFactoryTest(unittest.TestCase):
    """Test the SchemaFactory for generating Schemas.
    """
    def __init__(self, method):
        unittest.TestCase.__init__(self, method)

        # a cached schema bank, so we don't have to load it multiple times
        self.schema_bank = None

    def setUp(self):
        self.factory = Schema.SchemaFactory()

        self.test_file = os.path.join(os.getcwd(), "NeuralNetwork", "enolase.fasta")

        ambiguity_chars = {"G": "G",
                           "A": "A",
                           "T": "T",
                           "C": "C",
                           "R": "AG",
                           "*": "AGTC"}

        self.schema = Schema.Schema(ambiguity_chars)

    def test_easy_from_motifs(self):
        """Generating schema from a simple list of motifs.
        """
        motifs = {"GATCGAA": 20,
                  "GATCGAT": 15,
                  "GATTGAC": 25,
                  "TTTTTTT": 10}

        motif_bank = Pattern.PatternRepository(motifs)

        schema_bank = self.factory.from_motifs(motif_bank, .5, 2)
        if VERBOSE:
            print("\nSchemas:")
            for schema in schema_bank.get_all():
                print("%s: %s" % (schema, schema_bank.count(schema)))

    def test_hard_from_motifs(self):
        """Generating schema from a real life set of motifs.
        """
        schema_bank = self._load_schema_repository()

        if VERBOSE:
            print("\nSchemas:")
            for schema in schema_bank.get_top(5):
                print("%s: %s" % (schema, schema_bank.count(schema)))

    def _load_schema_repository(self):
        """Helper function to load a schema repository from a file.

        This also caches a schema bank, to prevent having to do this
        time consuming operation multiple times.
        """
        # if we already have a cached repository, return it
        if self.schema_bank is not None:
            return self.schema_bank

        # otherwise, we'll read in a new schema bank

        # read in the all of the motif records
        motif_handle = open(self.test_file, 'r')
        seq_records = list(SeqIO.parse(motif_handle, "fasta",
                                       alphabet=IUPAC.unambiguous_dna))
        motif_handle.close()

        # find motifs from the file
        motif_finder = Motif.MotifFinder()
        motif_size = 9

        motif_bank = motif_finder.find(seq_records, motif_size)

        schema_bank = self.factory.from_motifs(motif_bank, .1, 2)

        # cache the repository
        self.schema_bank = schema_bank

        return schema_bank

    def test_schema_representation(self):
        """Convert sequences into schema representations.
        """
        # get a set of schemas we want to code the sequence in
        schema_bank = self._load_schema_repository()
        top_schemas = schema_bank.get_top(25)
        schema_coder = Schema.SchemaCoder(top_schemas, self.schema)

        # get the sequences one at a time, and encode them
        fasta_handle = open(self.test_file, 'r')
        for seq_record in SeqIO.parse(fasta_handle, "fasta",
                                      alphabet=IUPAC.unambiguous_dna):
            schema_values = schema_coder.representation(seq_record.seq)
            if VERBOSE:
                print("Schema values: %s" % schema_values)
        fasta_handle.close()


# --- Tests for Signatures
class SignatureFinderTest(unittest.TestCase):
    """Test the ability to find signatures in a set of sequences.
    """
    def setUp(self):
        test_file = os.path.join('NeuralNetwork', 'enolase.fasta')

        self.test_records = []

        # load the records
        handle = open(test_file, 'r')
        self.test_records = list(SeqIO.parse(handle, "fasta",
                                             alphabet=IUPAC.unambiguous_dna))
        handle.close()

        self.sig_finder = Signature.SignatureFinder()

    def test_find(self):
        """Find signatures from sequence inputs.
        """
        repository = self.sig_finder.find(self.test_records, 6, 9)
        top_sig = repository.get_top(1)

        self.assertEqual(top_sig[0], ('TTGGAA', 'TGGAAA'))


class SignatureCoderTest(unittest.TestCase):
    """Test the ability to encode sequences as a set of signatures.
    """
    def setUp(self):
        signatures = [("GAC", "GAC"), ("AAA", "TTT"), ("CAA", "TTG")]

        self.coder = Signature.SignatureCoder(signatures, 9)

        self.test_seqs = [("GACAAAGACTTT", [1.0, 1.0, 0.0]),
                          ("CAAAGACGACTTTAAATTT", [0.5, 1.0, 0.0]),
                          ("AAATTTAAAGACTTTGAC", [1.0 / 3.0, 1.0, 0.0]),
                          ("GACGAC", [1.0, 0.0, 0.0]),
                          ("GACAAAAAAAAAGAC", [1.0, 0.0, 0.0]),
                          ("GACAAAAAAAAAAGAC", [0.0, 0.0, 0.0])]

    def test_representation(self):
        """Convert a sequence into its signature representation.
        """
        for seq_string, expected in self.test_seqs:
            test_seq = Seq(seq_string, IUPAC.unambiguous_dna)
            predicted = self.coder.representation(test_seq)

            self.assertEqual(predicted, expected,
                             "Non-expected representation %s for %s, wanted %s"
                             % (predicted, seq_string, expected))

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