#!/usr/bin/env python
"""Tests for an Organism in a Genetic Algorithm population.
"""
# standard library
import sys

# Biopython
from Bio import Alphabet
from Bio.Seq import MutableSeq

# PyUnit
import unittest

# local stuff
from Bio.GA import Organism

def run_tests(argv):
    test_suite = testing_suite()
    runner = unittest.TextTestRunner(sys.stdout, verbosity = 2)
    runner.run(test_suite)
    
def testing_suite():
    """Generate the set of tests.
    """
    test_suite = unittest.TestSuite()

    test_loader = unittest.TestLoader()
    test_loader.testMethodPrefix = 't_'
    tests = [CreatePopulationTest, OrganismTest]

    for test in tests:
        cur_suite = test_loader.loadTestsFromTestCase(test)
        test_suite.addTest(cur_suite)

    return test_suite

# -- utility functions
class TestAlphabet(Alphabet.Alphabet):
    """Simple alphabet for test purposes.
    """
    letters = ["1", "2", "3", "4"]

def genome_generator():
    """Generate a genome for testing purposes.
    """
    return MutableSeq("1234", TestAlphabet())

def fitness_calculator(genome):
    """Calculate fitness for testing purposes.
    """
    assert isinstance(genome, MutableSeq), "Expected MutableSeq for a genome."

    regular_seq = genome.toseq()
    return int(regular_seq.data)

class CreatePopulationTest(unittest.TestCase):
    """Tests for utility functions for creating populations.
    """
    def setUp(self):
        self.alphabet = TestAlphabet()

    def t_function_population(self):
        """Create a population using a function to generate genomes.
        """
        num_orgs = 10
        new_pop = Organism.function_population(genome_generator,
                                               num_orgs, fitness_calculator)

        assert len(new_pop) == num_orgs, "Expected %s organisms, got %s" \
               % (num_orgs, len(new_pops))

        for org in new_pop:
            assert isinstance(org, Organism.Organism), \
                   "Expected to get an organism, got %r" % org

            exp_fit = fitness_calculator(org.genome)
            assert org.fitness == exp_fit, \
                   "Expected fitness of %s, got %s" % (org.fitness, exp_fit)

    def t_random_population(self):
        """Create a population randomly from a alphabet.
        """
        num_orgs = 10
        genome_size = 5
        new_pop = Organism.random_population(self.alphabet, genome_size,
                                             num_orgs, fitness_calculator)

        assert len(new_pop) == num_orgs, "Expected %s organisms, got %s" \
               % (num_orgs, len(new_pops))

        for org in new_pop:
            assert isinstance(org, Organism.Organism), \
                   "Expected to get an organism, got %r" % org

            exp_fit = fitness_calculator(org.genome)
            assert org.fitness == exp_fit, \
                   "Expected fitness of %s, got %s" % (org.fitness, exp_fit)

            assert len(org.genome) == genome_size, \
                   "Expected genome size of %s, got %s" % (len(org.genome),
                                                           genome_size)

    def t_random_population_types(self):
        """Creating a random population with different types of alphabets.
        """
        class DoubleAlphabet:
            letters = [1.0, 2.0]

        class CharacterAlphabet:
            letters = ["a", "b"]

        class IntegerAlphabet:
            letters = [1, 2]

        def test_fitness(genome):
            return 2

        all_alphabets = [DoubleAlphabet(), CharacterAlphabet(),
                         IntegerAlphabet()]

        for alphabet in all_alphabets:
            new_pop = Organism.random_population(alphabet, 5, 10,
                                                 test_fitness)

class OrganismTest(unittest.TestCase):
    """Tests for an organism in a GA population.
    """
    def setUp(self):
        self.alphabet = TestAlphabet()
        self.genome = MutableSeq("1234", self.alphabet)
        self.organism = Organism.Organism(self.genome, fitness_calculator)

    def t_organism_basic(self):
        """Exercise basic organism functionality.
        """
        same_genome = MutableSeq("1234", self.alphabet)
        same_organism = Organism.Organism(same_genome, fitness_calculator)

        dif_genome = MutableSeq("1111", self.alphabet)
        dif_organism = Organism.Organism(dif_genome, fitness_calculator)

        assert self.organism == same_organism, \
               "Comparison doesn't work for identical organisms."

        assert self.organism != dif_organism, \
               "Comparison doesn't work for different organism."

    def t_organism_fitness(self):
        """Test the ability to deal with the fitness of the genome.
        """
        assert self.organism.fitness == 1234, \
               "Unexpected fitness %s" % self.organism.fitness
        
        new_genome = MutableSeq("1111", self.alphabet)
        self.organism.genome = new_genome
        self.organism.recalculate_fitness()

        assert self.organism.fitness == 1111, \
               "Unexpected fitness %s" % self.organism.fitness

    def t_organism_copy(self):
        """Test copying of organisms.
        """
        new_organism = self.organism.copy()

        new_organism.genome.append("1")

        assert new_organism.genome != self.organism.genome, \
               "Did not provide a copy of the organism."

    def t_provide_fitness(self):
        """Test that providing a pre-calculated fitness works.
        """
        def fitness_calc(genome):
            raise ValueError("Should not have been executed.")

        genome = self.organism.genome

        # make sure not supplying fitness works
        try:
            new_org = Organism.Organism(genome, fitness_calc)
            raise AssertionError("Did not calculate fitness when expected.")
        except ValueError:
            pass

        # make sure supplying fitness works
        new_org = Organism.Organism(genome, fitness_calc, 50)

if __name__ == "__main__":
    sys.exit(run_tests(sys.argv))