#!/usr/bin/env python
# -*- coding: utf-8 -*-
import string, sys, getopt
from random import Random

class WeaselException(Exception): pass

class WeaselGen:
  """Stores one generation of mutated strings as a sorted tuple list, with
  the most fit strings first [(score, 'string'), (score, 'string')...] the
  whole list is in self.all, the best string and best score are included 
  """
  def __init__(self, generation):
   self.gen = generation
   self.best_score, self.best_string = generation[0]

  def mean_fit(self):
    ''' Returns the mean fitness of a whole generation'''
    totalw = 0
    for score, st in self.all:
      totalw = totalw + score
    return totalw / float(len(self.all)) / float(len(self.best_score))


def make_initial(target):
  '''Make a random string, containing upper and lower case letters
  and the same length as the target sequence for the run
  '''
  return ''.join(Random().sample(string.letters + ' ', len(target)))

def mutate(template, mutationrate):
  '''Mutate a string at given rate without locking in chars that match the 
  target. The arguments are the template string followed by the probability 
  that each char is changed (there is a ~1/50 mutated character will actually 
  remain the same since the replacement is selected from the same pool of
  characters)
  '''
  l = list(template)
  for i in range(len(l)):
    if Random().random() <= mutationrate:
      l[i] = Random().sample(string.letters + " ", 1)[0]
    else: 
      continue
  return ''.join(l)
	
def mutate_locked(template, mutationrate, target):
  '''Mutate a string but if any any chars already match the target sequence
  leave them alone. Arguments are the template string, probability 
  that each char is changed and the target string 
  '''
  l = list(template) 
  for i in range(len(l)):
    if l[i] == target[i]:
      continue
    elif Random().random() <= mutationrate:
      l[i] = Random().sample(string.letters + " ", 1)[0]
    else: 
      continue
  return ''.join(l)

def make_gen(template, target, popsize, mutationrate, locked = False):
  '''Makes a new generation from a template string by adding new
  mutants (from mutate() above) then scores each against the target string
  and returns WeaselGen instance (defined above) 
	
  The arguments are the template string to base the new genetation, the 
  target string (to score offspring), the size of the population to make,
  the rate of mutation at each char and whether or not characters are 
  locked in place once they match the target (a boolean defaulted to
  false)
  '''
  generation = []
  genscores = []
  if locked == True:
    for i in range(int(popsize)):
      generation.append(mutate_locked(template, mutationrate, target))
  elif locked == False :
    for i in range(int(popsize)):
      generation.append(mutate(template, mutationrate))
  for string in generation:
    score = 0
    for i in range(len(string)):
      if string[i] == target[i]:
	score = score + 1
      else:
	continue
    stringscore = score, string
    genscores.append(stringscore)
  return WeaselGen(sorted(genscores, reverse = True))


def main():
  """
  Usage: weasel.py -[options] [popsize] [mutation rate] 

  Makes one run of Dawkins' weasel algorithm, displaying the best string 
  in each generation and the 'fitness' of that string (number of 
  characters matching the target string). 
	  Options
    -h or --help  display this help message
    -t target sequence (defaults to"Methinks it is like a weasel")
    -l use locking (once a character matchs the target  do not mutate)
    -o [file] log the fitness of the each generation to the specified file
	  Arguments
    popsize: the number of strings created in each generation  
    mutation rate: the probability of each letter mutating in each 
    geneartion (a floating point number between 0 and 1) 
	  Example
    python weasel.py -l -o outfile.csv 100 0.05
    (a run with locking, logged to 'outfile.csv', n=100 and u=0.05)
  """
  target = "Methinks it is like a weasel"
  locked = False
  logged = False
  opts, args = getopt.getopt(sys.argv[1:], "o:t:hl", ["help"]) 
  for o,a in opts:
    if o in ("-h", "--help"):
	print main.__doc__
	sys.exit(0)	
    elif o == "-t":
	target = a
    elif o == "-o":
      logged = True
      out_handle = a
    elif o == "-l":
      locked = True
  if len(args) != 2:
    raise WeaselException("takes 2 and only 2 arguments, see -h for help")
    sys.exit(1)
  popsize = float(args[0])
  mutrate = float(args[1])
  survivor = make_initial(target)	 
  print "evolving ", target, " from ", survivor
  print "locked is", locked
  generation = 0
  fitness = 0
  if logged == True:
    outfile = open(out_handle, 'w')
    print "logging to ", out_handle
  while fitness != len(target):
    generation = generation + 1
    population = make_gen(survivor, target, popsize, mutrate, locked)
    fitness = population.best_score
    survivor = population.best_string
    print generation, fitness, survivor
    if logged == True:
      outfile.write(str(fitness) + '\n')

if __name__ == '__main__':
	main()