from board import Board

import random
import numpy
import cPickle as pickle
import math

INF = float('inf')
RANDOM_WALKS = 15

class HeuristicMaker(object):

    def __init__(self):
        self.results = {}
        self.digestMap = {}
        try:
            with open('data.p', 'rb') as fp:
                self.results = pickle.load(fp)
        except:
            None

        size = 0
        for code in self.results:
            (scoreSum, count) = self.results[code]
            size += count
        print 'size: ' + str(size)
        print 'states: ' + str(len(self.results))
        
        for i in range(100000):
            print i
            if i % 10000 == 0: self.save()
            start = Board()
            self.randomWalk(start, 1)

        '''for digest in self.results:
            
            (s, c) = self.results[digest]
            if c > 5:
                print '--------'
                node = self.digestMap[digest]
                node.printBoard()
            
                h = float(s) / c
                print s
                print c
                print h'''
        
        self.save()
        print 'done'

    def save(self):
        with open('data.p', 'wb') as fp:
            pickle.dump(self.results, fp)
    
    def choseNextNode(self, node, t):
    
        nextNode = None
        maxValue = None
        for action in node.getLegalActions():
            next = node.getSuccessorState(action)  
            digest = next.digest()
            if not digest in self.results: return next
            (scoreSum, count) = self.results[digest]
            utility = float(scoreSum) / count
            if not node.isUserTurn():
                utility *= -1
            utility += 1
            bias = math.sqrt(4.0 * math.log(t) / count)
            value = utility + bias
            if nextNode == None or value > maxValue:
                nextNode = next
                maxValue = value
        return nextNode

    def randomWalk(self, node, t):
        if node.isTerminal():
            if node.isUserTurn(): return 1
            return -1
        actions = node.getLegalActions()
        if len(actions) == 0: return 0

        nodeCode = node.digest()
        nextNode = self.choseNextNode(node, t)         
        score = self.randomWalk(nextNode, t + 1)

        if not nodeCode in self.results:
            self.results[nodeCode] = (0,0)
            self.digestMap[nodeCode] = node
        (scoreSum, count) = self.results[nodeCode]
        count += 1
        scoreSum += score
        self.results[nodeCode] = (scoreSum, count)

        return score

