# -*- coding: utf-8 -*-
"""
Created on Tue Apr 12 13:15:12 2016

@author: Samir Kanaan
"""


# This class encloses ten oil wells, each with an extract() operation that
# returns the amount of barrels produced in that extraction.

import operator
import random

import numpy

from deap import base
from deap import benchmarks
from deap import creator
from deap import tools

class OilWells(object):
    # Each oil well is a function y = ax**2+bx+c
    def __init__(self):
        self.a = [ 0.005, 0.001, -0.1, 0.002,    0, -0.05,  -0.1, 0.01, 0.001,  0.02]
        self.b = [  -1.5, -0.59, 0.05,  -0.2, -0.1,  0.05, -0.05, -1.6, 0.001,  -1.8]
        self.c = [    40,    28,   42,    10,   12,    30,    45,   38,     3,    22]

        # Current extraction number in each well
        self.x = [0]*self.numberOfWells()

    # Simply return the number of wells
    def numberOfWells(self):
        return len(self.a)

    # Reset the state of the wells to allow a new simulation
    def reset(self):
        self.x = [0]*self.numberOfWells()


    # Make an extraction from a given well (0..9). Min is 0 barrels
    def extract(self, well):
        try:
            x = self.x[well]
            y = self.a[well]*x**2 + self.b[well]*x + self.c[well]
            self.x[well] += 1
            return max(0,round(y))
        except:
            print('ERROR: invalid well number(%d)' % well)



creator.create("OilWellsMax", OilWells, weights=(1.0,), well = None, barrels = None)
creator.create("Well", list, well=creator.OilWellsMax)


def generate( well):
    part = creator.OilWellsMax()
    part.barrels = part.extract( well)
    part.well = well
    return part


toolbox = base.Toolbox()
toolbox.register("evaluate", generate, well = random.randint(0,9) )
toolbox.register("individual", tools.initRepeat, creator.Well, toolbox.evaluate,  n = 100)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)




def main():
    pop = toolbox.population(n = 10)

    stats = tools.Statistics(lambda ind: ind.well.barrels)
    stats.register("max", numpy.max)


    logbook = tools.Logbook()
    logbook.header = ["gen", "evals"] + stats.fields

    GEN = 1000
    best = None

    for g in range(GEN):
        for part in pop:
            #part.barrels = toolbox.evaluate(part).barrels
            if not part.best or part.best.well < part.well:
                part.best = creator.Well(part)
                part.best.well.barrels = part.well.barrels
            if not best or best.well < part.well:
                best = creator.Well(part)
                best.well.barrels = part.well.barrels
        for part in pop:
            toolbox.update(part, best)

        # Gather all the fitnesses in one list and print the stats
        logbook.record(gen=g, evals=len(pop), **stats.compile(pop))
        print(logbook.stream)

    return pop, logbook, best

main()