from multiprocessing.pool import ThreadPool
import multiprocessing as mpc
from time import time as timer
import math

import array
import numpy
import random

from deap import base
from deap import creator
from deap import tools
from deap import algorithms

import argparse
import os


parser = argparse.ArgumentParser(description='Executes GA')
parser.add_argument('--gens',
                    help='Number of generations to run.',
                    type=int,
                    default=2)
parser.add_argument('--pop_size',
                    help='Population size per generation.',
                    type=int,
                    default=5)
parser.add_argument('--save_frames',
                    help='Write frames to file.',
                    action='store_true')

args = parser.parse_args()
from bouncing_balls import BouncyBalls

# Draw two random lines with the intent of maximizing the number of balls
# on the screen
def eval2DPhysics(individual, last=False): #TBD - do something with indiv!
  if not last:
    os.environ['SDL_VIDEODRIVER'] = 'dummy' # Run pygame headless
  else:
    os.environ['SDL_VIDEODRIVER'] = 'x11'   # Run pygame headfull

  l1 = (individual[0],individual[1],individual[2],individual[3])
  l2 = (individual[4],individual[5],individual[6],individual[7])
  game = BouncyBalls(l1,l2)
  return game.run(),

creator.create("FitnessMax", base.Fitness, weights=(1.0,))
creator.create("Individual", list, fitness=creator.FitnessMax)

toolbox = base.Toolbox()

toolbox.register("attr_int", random.randint, 0, 600) #make pair if not square resolution
toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_int, 8)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)


pool = mpc.Pool()
toolbox.register("map", pool.map)
toolbox.register("evaluate", eval2DPhysics)
toolbox.register("mate", tools.cxTwoPoint)
toolbox.register("mutate", tools.mutUniformInt, low=0, up=600, indpb=0.2)
#toolbox.register("mutate", tools.mutGaussian, mu=0.0, sigma=0.2, indpb=0.2)
toolbox.register("select", tools.selTournament, tournsize=3)

def main(gens, pop_size, save_frames=False):
  pop = toolbox.population(n=pop_size)
  CXPB, MUTPB = 0.5, 0.2

  fitnesses = toolbox.map(toolbox.evaluate, pop)

  # Copy fitnesses to population
  for ind, fit in zip(pop, fitnesses):
    ind.fitness.values = fit
  fits = [ind.fitness.values[0] for ind in pop]

  gen = 0
  while gen < gens:
    gen += 1
    print("Generation %d" % gen)

    offspring = toolbox.select(pop, len(pop))
    offspring = list(map(toolbox.clone, offspring))

    # Crossover
    for c1, c2 in zip(offspring[::2], offspring[1::2]):
      if random.random() < CXPB:
        toolbox.mate(c1, c2)
        del c1.fitness.values
        del c2.fitness.values

    # Mutation
    for mutant in offspring:
      if random.random() < MUTPB:
        toolbox.mutate(mutant)
        for m in mutant:
            if m < 0.0: m = 0.0
            if m > 600.0:  m = 600.0
        del mutant.fitness.values

    invalid_ind = [ind for ind in offspring if not ind.fitness.valid]

    # Evaluation
    fitnesses = toolbox.map(toolbox.evaluate, invalid_ind)

    for ind, fit in zip(invalid_ind, fitnesses):
      ind.fitness.values = fit

    pop[:] = offspring
    fits = [ind.fitness.values[0] for ind in pop]

    print("%d invalid indices evaluated" % len(invalid_ind))

    # Print statistics
    length = len(pop)
    mean = sum(fits) / length
    sum2 = sum(x*x for x in fits)
    std = abs(sum2/length - mean**2)**0.5

    print("* Min: %s" % min(fits))
    print("* Max: %s" % max(fits))
    print("* Avg: %s" % mean)
    print("* Std: %s" % std)
    best = tools.selBest(pop, 1)[0]
    print("* Best: %s, %s" % (best, best.fitness.values))#, eval2DPhysics(best, True)))

  print("Done.")
  best_ind = tools.selBest(pop, 1)[0]
  print("Best individual: %s, %s, %s" % (best_ind, best_ind.fitness.values, eval2DPhysics(best_ind, True)))

if __name__ == '__main__':
  assert sys.version_info >= (3, 7), "Requires Python 3.7+"
  main(args.gens, args.pop_size, args.save_frames)