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

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


# local stuff
from Bio.GA import Organism


# -- 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(str(regular_seq))


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

    def test_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 test_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 test_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 test_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 str(self.organism) == str(same_organism), \
               "Comparison doesn't work for identical organisms."

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

    def test_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 test_organism_copy(self):
        """Test copying of organisms.
        """
        new_organism = self.organism.copy()

        new_organism.genome.append("1")

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

    def test_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__":
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)
