import collections
import math
import itertools
import time
from PIL import ImageGrab

import win32api
import win32con

INF = 1000000


def PressKey(hexKeyCode):
    win32api.keybd_event(hexKeyCode, 0, 0, 0);

def ReleaseKey(hexKeyCode):
    win32api.keybd_event(hexKeyCode, 0, win32con.KEYEVENTF_KEYUP, 0)

# Actuals Functions


# For each row define the shifts
shiftLeft = [None for i in range(0xFFFF)]
shiftRight = [None for i in range(0xFFFF)]
shiftUp = {}
shiftDown = {}

def genShiftResults():
    global shiftLeft, shiftRight, shiftUp, shiftDown

    def row2Int(r, shift, reverse=False):
        if reverse:
            r = r[::-1]
        return (r[0] << (3*shift)) + (r[1] << (2*shift)) + (r[2] << shift) + r[3]

    possible = list(range(0, 0xF))
    for a, b, c, d in itertools.product(possible, repeat=4):
        rowList = [a,b,c,d]
        newList = [0,0,0,0]

        i = 0
        last = -1
        for j in range(4):
            if rowList[j] == last:
                newList[i-1] = rowList[j] + 1
                last = -1
            elif rowList[j] > 0:
                newList[i] = rowList[j]
                last = rowList[j]
                i += 1

        shiftRight[row2Int(rowList, 4, False)] = row2Int(newList, 4, False)
        shiftLeft[row2Int(rowList, 4, True)]   = row2Int(newList, 4, True)

        shiftUp[row2Int(rowList, 16, True)]   = row2Int(newList, 16, True)
        shiftDown[row2Int(rowList, 16, False)]  = row2Int(newList, 16, False)



def getValue(state, x, y):
    t = (state >> (4 * (4 * y + x))) & 0xF
    assert 0 <= t <= 0xF
    return t

def modifiedState(state, x, y, val):
    assert 0 <= val <= 0xF
    maskB = 4 * (4 * y + x)
    state ^= (state & (0xF << maskB))
    state |= val << maskB
    return state

def printState(state):
    for y in range(4):
        for x in range(4):
            item = getValue(state, x, y)
            string = "{0:>4} ".format(2 ** item if item > 0 else "_")
            print(string, end="")
        print()
    print()

def children(state):
    # TODO: support 4's appearing and relative probability
    tempState = state
    for y in range(4):
        for x in range(4):
            if (tempState & 0xF) == 0:
                yield modifiedState(state, x, y, 2)
            tempState >>= 4

def shiftAbout(tempState, shiftArray, andMask, shift):
    new = 0
    for column in range(4):
        new |= shiftArray[tempState & andMask] << shift * column
        tempState >>= shift
    return new

def move(state, direction, gravity):
    global shiftLeft, shiftRight, shiftUp, shiftDown
    global columnMask
    # direction is 0,1
    # gravity is -1,1

    if direction == 0 and gravity == 1:
        moved = shiftAbout(state, shiftLeft, 0xFFFF, 16)

    elif direction == 0 and gravity == -1:
        moved = shiftAbout(state, shiftRight, 0xFFFF, 16)

    elif direction == 1 and gravity == 1:
        moved = shiftAbout(state, shiftUp, 0xF000F000F000F, 4)

    elif direction == 1 and gravity == -1:
        moved = shiftAbout(state, shiftDown, 0xF000F000F000F, 4)

    return moved != state, moved


# Aims to maximize heuristic
#sqrtPieceHeuristic(state):


# -150 to 0 range
def heuristicSqrtPieces(state):
    heur = 0
    for y in range(4):
        for x in range(4):
            v = getValue(state, x, y)
            if v > 0:
                heur -= math.sqrt(1 << v)
    return heur



# -100 to 0 range
def heuristicNotCentered(state):
    maxTile = max(getValue(state, x, y) for y in range(4) for x in range(4))

    heur = 0
    for y in range(1,3):
        for x in range(1,3):
            v = getValue(state, x, y)
            # try to keep four largest on outside
            if v > (maxTile - 4):
                heur -= 10 * v
    return heur


# 0 to 15 (higher is better)
def heuristicBlanks(state):
    heur = 0
    for y in range(4):
        for x in range(4):
            piece = getValue(state, x, y)
            heur += piece == 0
    return heur


calls = 0
evaled = 0
saved = {}
def BeamAStarAlphaBeta(state, depth = 3):
    global calls, evaled, saved

    calls += 1

    key = (state, depth)
    temp = saved.get(key, None)
    if temp:
        return temp

    evaled += 1

    if depth == 0:
        heur = heuristicSqrtPieces(state)
        heur += heuristicNotCentered(state) / 2
        heur += 10 * heuristicBlanks(state)
        return (heur, state, "")

    best = -INF
    bestAdvice = None
    bestResult = None

    dirs = [("Left", 0, 1), ("Up", 1, 1), ("Right", 0, -1), ("Down", 1, -1)]
    for name, direction, gravity in dirs:
        legal, result = move(state, direction, gravity)

        if legal:
            sumChildren = 0
            countChildren = 0
            for child in children(result):
                fitness, r, a = BeamAStarAlphaBeta(child, depth-1)
                sumChildren += fitness
                countChildren += 1

            if countChildren > 0:
                avgChild = sumChildren / float(countChildren)
            else:
                avgChild = -INF

            heuristic = avgChild

            if heuristic > best:
                best = heuristic
                bestAdvice = name
                bestResult = result


    result = (best, bestResult, bestAdvice)
    saved[key] = result
    return result


def main(activeTilesFromScreenCapture):
    global calls, evaled, saved

    while True:
        time.sleep(0.40)
        state = activeTilesFromScreenCapture()
        printState(state)

        tileToDepth = {512:4, 1024:5, 2048:5, 4096:5, 8192:6}
        maxTile = max(getValue(state, x, y) for x in range(4) for y in range(4))
        freeTiles = heuristicBlanks(state)

        depth = tileToDepth.get(maxTile, 3) + (freeTiles < 5)

        calls = 0
        evaled = 0
        saved = {}
        T0 = time.time()
        fitness, result, advice = BeamAStarAlphaBeta(state, depth)
        T1 = time.time() + 0.00001

        print ("advice: {} : {:.3f} (depth: {} => {} calls {} evaled) ({:.1f} positons / s)".format(
            advice, fitness, depth, calls, evaled, evaled / (T1 - T0)))
        print ()
        print ()

        keyMapping = {
            "Left" : 0x025,
            "Right" : 0x027,
            "Up" : 0x026,
            "Down" : 0x028
        }

        key = keyMapping.get(advice)
        PressKey(key)
        time.sleep(0.05)
        ReleaseKey(key)



def calibrateScreen():
    # TODO: anything but this just please

    # play with these numbers till you get the entire board
    cords = [225,352,720,852]
    img = ImageGrab.grab(cords)

    width = cords[2] - cords[0]
    height = cords[3] - cords[1]
    #print (width, height)

    # Squirrel math that hopefully works for you
    average = (width + height) // 2
    margins = int(average * 0.15)
    margin = margins // 4
    squareSize = (average - margins) // 4
    #print (squareSize)

    squareOffset = margin + squareSize // 12

    tileCords = [[None for i in range(4)] for j in range(4)]

    for x in range(12):
        for y in range(10):
            purple = 255*255*255 + 255*255
            img.putpixel((squareSize+margin+x, squareSize+margin+y), purple)

    for xSquare in range(4):
        for ySquare in range(4):
            xStart = squareOffset + (squareSize + margin) * xSquare
            yStart = squareOffset + (squareSize + margin) * ySquare
            tileCords[xSquare][ySquare] = (xStart,yStart)

#            for i in range(10):
#                for j in range(10):
#                    img.putpixel((xStart+i, yStart+j), 255 * 255)

 #   img.show()

    def activeTilesFromScreenCapture():
        img = ImageGrab.grab(cords)
        state = 0
        for x in range(4):
            for y in range(4):
                tileType = {
                    (238, 228, 218) : 1,
                    (237, 224, 200) : 2,
                    (242, 177, 121) : 3,
                    (245, 149, 99)  : 4,
                    (246, 124, 95)  : 5,
                    (246, 94, 59)   : 6,
                    (237, 207, 114) : 7,
                    (237, 204, 97)  : 8,
                    (237, 200, 80)  : 9,
                    (237, 197, 63)  : 10,
                    (237, 194, 46)  : 11
                }
                pixel = img.getpixel((tileCords[x][y]))
                tType = tileType.get(pixel, -1)
                if  (abs(pixel[0] - 204) +
                    abs(pixel[1] - 192) +
                    abs(pixel[2] - 179)) < 10:
                    tType = 0

                if tType < 0:
                    print (pixel, tType, "??? Pausing for 5s")
                    time.sleep(5)
                    return activeTilesFromScreenCapture()

                state = modifiedState(state, x, y, tType)
        return state

    return activeTilesFromScreenCapture


if __name__ == "__main__":
    genShiftResults()

    activeTilesFromScreenCapture = calibrateScreen()
    start = activeTilesFromScreenCapture()
    print (start)

    '''
    dirs = [("Left", 0, 1), ("Up", 1, 1), ("Right", 0, -1), ("Down", 1, -1)]
    for name, direction, gravity in dirs:
        legal, result = move(start, direction, gravity)
        print (name, legal)
        printState(result)

#'''
    main(activeTilesFromScreenCapture)



# v4 was 15000 evaled / second
# v5 was 17000 evaled / second