from random import random, randint, choice
from copy import deepcopy
from math import log

class fwrapper:
    def __init__(self, function, childcount, name):
        self.function = function
        self.childcount = childcount
        self.name = name

class node:
    def __init__(self, fw, children):
        self.function = fw.function
        self.name = fw.name
        self.children = children

    def evaluate(self, inp):
        results = [n.evaluate(inp) for n in self.children]
        return self.function(results)

    def display(self, indent = 0):
        # output as lisp code
        print((' ' * indent) + '(' + self.name)
        for c in self.children:
            c.display(indent + 1)
        print(')')

class paramnode:
    def __init__(self, idx):
        self.idx = idx

    def evaluate(self, inp):
        return inp[self.idx]

    def display(self, indent = 0):
        # output as lisp code
        print('%sp%d' % (' ' * indent, self.idx))

class constnode:
    def __init__(self, v):
        self.v = v

    def evaluate(self, inp):
        return self.v

    def display(self, indent = 0):
        # output as lisp code
        print('%s%d' % (' ' * indent, self.v))

def makerandomtree(pc, maxdepth = 10, fpr = 0.5, ppr = 0.6):
    flist = [
        fwrapper(lambda l:l[1] if l[0] != 0 else l[2], 3, 'if'),
        fwrapper(lambda l:1 if l[0] > l[1] else 0, 2, '>'),
        fwrapper(lambda l:1 if l[0] < l[1] else 0, 2, '<'),
        fwrapper(lambda l:l[0] + l[1], 2, '+'),
        fwrapper(lambda l:l[0] - l[1], 2, '-'),
        fwrapper(lambda l:l[0] * l[1], 2, '*'),
        # ignore division as it brings too much trouble
    ]

    if random() < fpr and maxdepth > 0:
        f = choice(flist)
        children = [makerandomtree(pc, maxdepth - 1, fpr, ppr)
                    for i in range(f.childcount)]
        return node(f, children)
    elif random() < ppr:
        return paramnode(randint(0, pc - 1))
    else:
        return constnode(randint(0, 10))

def hiddenfunction(x, y):
    # the truth of the nature
    return x**2 + y**2 + 1

def buildhiddenset():
    # observed phenomenons
    rows = []
    for i in range(1000):
        x = randint(0, 100)
        y = randint(0, 100)
        rows.append([x, y, hiddenfunction(x, y)])
    return rows

def scorefunction(tree, s):
    dif = 0
    for data in s:
        v = tree.evaluate([data[0], data[1]])
        dif += abs(v - data[2])
    return dif

def mutate(tree, pc, probchange = 0.3):
    if random() < probchange:
        return makerandomtree(pc)
    else:
        result = deepcopy(tree)
        if isinstance(tree, node):
            result.children = [mutate(c, pc, probchange)
                               for c in tree.children]
        return result

def crossover(tree1, tree2, probswap = 0.7, top = 1):
    if random() < probswap and not top:
        return deepcopy(tree2)
    else:
        result = deepcopy(tree1)
        if isinstance(tree1, node) and isinstance(tree2, node):
            result.children = [crossover(c, choice(tree2.children), probswap, 0)
                               for c in tree1.children]
        return result

def evolve(pc, popsize, rankfunction, maxgen = 500, mutationrate = 0.1, breedingrate = 0.4, pexp = 0.7, pnew = 0.05):
    def selectindex():
        # equivalent to log_pexp^random(), function value from zero to infinity
        return int(log(random())/log(pexp))

    population = [makerandomtree(pc) for i in range(popsize)]
    for i in range(maxgen):
        scores = rankfunction(population)
        print(scores[0][0])
        if scores[0][0] == 0:
            break

        # at least we need a couple
        newpop = [scores[0][1], scores[1][1]]
        while len(newpop) < popsize:
            if random() > pnew:
                # breed
                newpop.append(
                    mutate(
                        crossover(
                            scores[min(selectindex(), len(scores) - 1)][1],
                            scores[min(selectindex(), len(scores) - 1)][1],
                            probswap = breedingrate),
                        pc, probchange = mutationrate))
            else:
                # make new out of void
                newpop.append(makerandomtree(pc))
        population = newpop
    scores[0][1].display()
    return scores[0][1]

def getrankfunction(dataset):
    def rankfunction(population):
        scores = [(scorefunction(t, dataset), t)
                  for t in population]
        scores = sorted(scores, key=lambda x: x[0])
        return scores
    return rankfunction

evolve(2, 300, getrankfunction(buildhiddenset()), 500, 0.3, 0.3, 0.2, 0.2)