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

class fWrapper:
    def __init__(self,function,child_count,name,params,type):
        self.function = function
        self.child_count = child_count
        self.name = name
        self.params = params
        self.type = type

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

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

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

class paramNode:
    def __init__(self,index):
        self.index = index
        self.name = "p"+str(index)
        self.type = 'int'

    def evaluate(self,input):
        return input[self.index]

    def display(self,indent=0):
        print '%sp%d' % (' '*indent,self.index)

class constNode:
    def __init__(self,v):
        self.v = v
        self.name = str(v)
        self.type = 'string' if isinstance(v,str) else 'int'

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

    def display(self,indent=0):
        if isinstance(self.v,str):
            print '%s%s' % (' '*indent,self.v)
        else:
            print '%s%d' % (' '*indent,self.v)

def div_func(l):
    return l[0]//l[1] if l[1] != 0 else l[0]

def mod_func(l):
    return l[0]%l[1] if l[1] != 0 else l[0]

def if_func(l):
    if isinstance(l[0],int) and l[0]>0:
        return l[1] 
    elif isinstance(l[0],str) and len(l[0])>0:
        return l[1]
    else: return l[2]

def is_greater(l):
    if l[0] > l[1]: return 1
    else: return 0

def fizz_func(l):
    return "fizz"

def buzz_func(l):
    return "buzz"

def concat(l):
    return l[0]+l[1]

addw = fWrapper(lambda l:l[0]+l[1],2,'add',['int','int'],'int')
subw = fWrapper(lambda l:l[0]-l[1],2,'sub',['int','int'],'int')
mulw = fWrapper(lambda l:l[0]*l[1],2,'mul',['int','int'],'int')
divw = fWrapper(div_func,2,'div',['int','int'],'int')
modw = fWrapper(mod_func,2,'mod',['int','int'],'int')

ifw = fWrapper(if_func,3,'if',['object','object','object'],'object')
gtw = fWrapper(is_greater,2,'is_greater',['int','int'],'int')

fizzw = fWrapper(fizz_func,0,'fizz',[],'string')
buzzw = fWrapper(buzz_func,0,'buzz',[],'string')
concatw = fWrapper(concat,2,'concat',['string','string'],'string') 

flist = {'int':[divw,modw],'string':[fizzw,buzzw,concatw]}
flist['object'] = [ifw] + flist['int'] + flist['string']

def randstr(n):
    alphabets = string.digits + string.letters 
    return ''.join(choice(alphabets) for i in xrange(n))

def make_random_tree(pc,max_depth=4,fpr=0.5,ppr=0.6,type='object'):
    if random() < fpr and max_depth > 0:
        f = choice(flist[type])
        children = [make_random_tree(pc,max_depth-1,fpr,ppr,type=t) for t in f.params]
        return node(f,children)
    elif random() < ppr and type == 'int':
        return paramNode(randint(0,pc-1))
    elif type == 'object':
        return constNode(choice(['','fizz','buzz','fizzbuzz',randint(0,100)]));
    elif type == 'string':
        return constNode(choice(['','fizz','buzz','fizzbuzz']))
    else:
        return constNode(randint(0,100))

def hidden_function(x):
    if x % 15 == 0:
        return 'fizzbuzz'
    elif x % 3 == 0:
        return 'fizz'
    elif x % 5 == 0:
        return 'buzz'
    else:
        return ''

def build_hidden_set():
    rows = []
    for i in range(1,101):
        rows.append([i,hidden_function(i)])
    return rows

def score_function(tree,s):
    diff = 0
    for data in s:
        v = tree.evaluate([data[0]])
        if isinstance(v,str):
            if v != data[1]:
                diff += 1
        else:
            diff += 20
    return diff

def mutate(t,pc,prob_change=0.1,type='string'):
    if random() < prob_change:
        return make_random_tree(pc,type=type)
    else:
        result = deepcopy(t)
        if hasattr(t,"children"):
            result.children = [mutate(c,pc,prob_change,type=c.type) for c in t.children]
        return result
    return deepcopy(t)

def crossover(t1,t2,prob_swap=0.7,top=1):
    def choose_from_same_type(child,tree):
        nodes_with_same_type = filter(lambda n: n.type == child.type, tree.children)
        if len(nodes_with_same_type) == 0:
            return None
        else:
            return choice(nodes_with_same_type)

    if t2 == None:
        return deepcopy(t1)
    if random() < prob_swap and not top:
        return deepcopy(t2)
    else:
        result = deepcopy(t1)
        if hasattr(t1,'children') and hasattr(t2,'children'):
            result.children = [crossover(c,choose_from_same_type(c,t2),prob_swap,0) for c in t1.children]
        return result
    return deepcopy(t1)

def evolve(pc,popsize,rank_function,maxgen=500,mutation_rate=0.1,breeding_rate=0.4,pexp=0.7,pnew=0.05):
    def select_index():
        return int(log(random()) / log(pexp))

    population = [make_random_tree(pc,type='object') for i in range(popsize)]
    for i in range (maxgen):
        scores = rank_function(population)
        print i,scores[0][0]
        if i % 10 == 0:
            scores[0][1].display()
        if scores[0][0] == 0 : break

        newpop = [scores[0][1],scores[1][1]]

        while len(newpop) < popsize:
            if random() > pnew:
                newpop.append(mutate(crossover(scores[select_index()][1],scores[select_index()][1],prob_swap=breeding_rate),pc,prob_change=mutation_rate,type='object'))
            else:
                newpop.append(make_random_tree(pc,type='object'))
        population = newpop
    scores[0][1].display()
    write_jpeg(scores[0][1])

    return scores[0][1]

def get_rank_function(data_set):
    def rank_function(population):
        scores = [(score_function(t,data_set),t) for t in population]
        scores.sort()
        return scores
    return rank_function

def write_jpeg(tree):
    stack = [tree]
    g = pydot.Dot()
    root = True
    while len(stack) > 0:
        node = stack.pop()
        parent_node = pydot.Node(node.name+"_"+str(id(node)))
        parent_node.set_label(node.name)
        if root:
            g.add_node(parent_node)
            root = False
        if hasattr(node,"children"):
            for child in node.children:
                stack.append(child)
                child_node = pydot.Node(child.name+"_"+str(id(child)))
                child_node.set_label(child.name)
                g.add_node(child_node)
                g.add_edge(pydot.Edge(parent_node,child_node))
    g.write_jpeg('tree.jpg',prog='dot')
    
rf = get_rank_function(build_hidden_set())
evolve(1,100,rf,maxgen=5000,mutation_rate=0.2,breeding_rate=0.3,pexp=0.65,pnew=0.25)