#! /usr/bin/env python

"""
           A simple puzzle game.

Push pieces around to reassemble the image.
The image comes from a file that you specify,
which is divided in N times M blocks, shuffled
around randomly.  Use the BACKSPACE key to undo
your moves.

Usage:

  puzzle.py [-b] imagefilename [NxM] [PxQ]
  puzzle.py [-b] imagefilename savedboardname

  NxM: number of divisions, defaults to 3x3
  PxQ: board size, defaults to 3 more than NxM
  -b : flag to forbid blocks from touching the
         border of the board.

In the second form, the layout 'savedboardname'
is loaded from the file boards.py; try 'B11' or
'current'.
"""

import pygame
import pygame.event, pygame.key, pygame.transform
from pygame.locals import *
import random, time


DIRS = [(0, -1), (-1, 0), (0, 1), (1, 0)]

KEY2DIR = {
    K_LEFT:  (-1, 0),
    K_RIGHT: ( 1, 0),
    K_UP:    ( 0,-1),
    K_DOWN:  ( 0, 1),
    }
UNDO_KEY = K_BACKSPACE       # undo any number of moves
QUIT_KEY = K_ESCAPE          # quit

CHEAT_MOUSE_CLICK = False
CHEAT_SOLUTION_DB = False

#FULLSCREEN_RES = None
FULLSCREEN_RES = (1024, 768)


class Game:
    COLOR = {' ': (0, 0, 0),
             '@': (192, 0, 0),
             '&': (192, 0, 0),
             '#': (128, 128, 0),
             '*': (96, 96, 0),
             '.': (64, 0, 0),
             }

    def __init__(self, bwidth, bheight, cw, ch, img):
        self.bwidth = bwidth
        self.bheight = bheight
        self.cw = cw
        self.ch = ch
        pygame.display.init()
        pygame.key.set_repeat(400, 30)
        if FULLSCREEN_RES:
            winstyle = NOFRAME
            width, height = FULLSCREEN_RES
        else:
            winstyle = 0
            width = int((bwidth-1.6)*cw)
            height = int((bheight-1.6)*ch)
        self.screen = pygame.display.set_mode((width, height), winstyle)
        self.img = img.convert()
        self.deltax = (cw*bwidth - width) // 2
        self.deltay = (ch*bheight - height) // 2
        self.setup_empty_board()

    def setup_empty_board(self):
        width, height = self.screen.get_size()
        bwidth, bheight = self.bwidth, self.bheight
        self.board = {}
        self.history = None
        xmin, ymin = self.screen2xy((0, 0))
        xmax, ymax = self.screen2xy((width-1, height-1))
        for y in range(ymin, ymax+1):
            for x in range(xmin, xmax+1):
                if not (0 <= x < bwidth and 0 <= y < bheight):
                    self.set(x, y, '*')
        self.board.clear()
        for x in range(bwidth):
            for y in range(bheight):
                self.set(x, y, ' ')
        for x in range(bwidth):
            self.set(x, 0, '#')
            self.set(x, self.bheight-1, '#')
        for y in range(bheight):
            self.set(0, y, '#')
            self.set(self.bwidth-1, y, '#')

        self.canplace = lambda (px, py), (bx, by): True
        self.startpos = (1, 1)

    def screen2xy(self, (mx, my)):
        cw = self.cw
        ch = self.ch
        mx += self.deltax
        my += self.deltay
        return mx/cw, my/ch

    def xy2screen(self, (x, y)):
        cw = self.cw
        ch = self.ch
        x1 = x*cw - self.deltax
        y1 = y*ch - self.deltay
        return x1, y1

    def set(self, x, y, piece, rec=True):
        cw = self.cw
        ch = self.ch
        x1, y1 = self.xy2screen((x, y))
        s = self.screen
        if self.history is not None and (x, y) not in self.history:
            self.history[x, y] = self.board[x, y]
        self.board[x, y] = piece
        border = 15
        if isinstance(piece, tuple):
            r, g, b = 160, 80, 128
            px, py = piece
            s.blit(self.img, (x1, y1), (px*cw, py*ch, cw, ch))
            if self.board[x, y-1] == (px, py-1): border &= ~1
            if self.board[x-1, y] == (px-1, py): border &= ~2
            if self.board[x, y+1] == (px, py+1): border &= ~4
            if self.board[x+1, y] == (px+1, py): border &= ~8
        else:
            r, g, b = self.COLOR[piece]
            s.fill((r, g, b), (x1, y1, cw, ch))
            if piece == ' ':
                return
        gap = 2 * bool(border&8)
        if border&1: s.fill((r*5//4, g*5//4, b*5//4), (x1, y1, cw-gap, 2))   #top
        gap = 2 * bool(border&4)
        if border&2: s.fill((r*5//4, g*5//4, b*5//4), (x1, y1, 2, ch-gap))   #lft
        gap = 2 * bool(border&2)
        if border&4: s.fill((r//2, g//2, b//2), (x1+gap, y1+ch-2, cw-gap, 2))#bot
        gap = 2 * bool(border&1)
        if border&8: s.fill((r//2, g//2, b//2), (x1+cw-2, y1+gap, 2, ch-gap))#rgt
        if rec:
            do = []
            if not (border&1): do.append((x,y-1))
            if not (border&2): do.append((x-1,y))
            if not (border&4): do.append((x,y+1))
            if not (border&8): do.append((x+1,y))
            for x, y in do:
                self.set(x, y, self.board[x, y], rec=False)

    def flip(self):
        pygame.display.flip()

    def insert_piece(self, piece, (piecesw, piecesh), no_border=False):
        xstart = 1 + (piece[0] == piecesw-1 or no_border)
        xstop = self.bwidth-1 - (piece[0] == 0 or no_border)
        ystart = 1 + (piece[1] == piecesh-1 or no_border)
        ystop = self.bheight-1 - (piece[1] == 0 or no_border)
        k = 0.0
        r = 0
        while True:
            r += 1
            if r > 20000:
                raise PlacementError
            nx = random.randrange(xstart, xstop)
            ny = random.randrange(ystart, ystop)
            if self.board[nx, ny] != ' ' or not self.canplace(piece, (nx, ny)):
                continue
            count = 0
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    count += self.board[nx+dx, ny+dy] not in (' ', '*', '.')
            if count <= k:
                break
            k += 0.02
        self.set(nx, ny, piece)

    def bigpiece(self, x0, y0):
        result = {}
        if not isinstance(self.board[x0, y0], tuple):
            return result
        px0, py0 = self.board[x0, y0]
        pending = [(x0, y0)]
        for x, y in pending:
            if (x, y) in result:
                continue
            if self.board[x, y] != (px0 + x-x0, py0 + y-y0):
                continue
            result[x, y] = True
            for dx, dy in DIRS:
                pending.append((x+dx, y+dy))
        return result

    def finished(self, nb_pieces):
        for y in range(self.bheight):
            for x in range(self.bwidth):
                bp = self.bigpiece(x, y)
                if len(bp) == nb_pieces:
                    return (x, y)
                if bp:
                    break
        return None


def nextevent():
    e = pygame.event.wait()
    if e.type == QUIT or (e.type == KEYDOWN and e.key == QUIT_KEY):
        raise SystemExit
    return e

def peeknextevent():
    e = pygame.event.poll()
    if e.type == QUIT or (e.type == KEYDOWN and e.key == QUIT_KEY):
        raise SystemExit
    return e

class PlacementError(Exception):
    pass

bk2player = {' ': '@',
             '.': '&'}
player2bk = {'@': ' ',
             '&': '.'}

def puzzle((piecesw, piecesh), (boardw, boardh), img, largeimg,
           bumpyborder=False, preset=None):
    imgw, imgh = img.get_size()
    game = Game(boardw+2, boardh+2, imgw//piecesw, imgh//piecesh, img)
    nb_pieces = piecesw*piecesh
    pieces_max = (piecesw, piecesh)

    def saveboard():
        getboarddict()[boardname] = boardinfo, initialboard

    if preset is None:
        while True:
            lst = [(x, y) for x in range(piecesw) for y in range(piecesh)]
            random.shuffle(lst)
            try:
                for piece in lst:
                    game.insert_piece(piece, pieces_max, bumpyborder)
            except PlacementError:
                peeknextevent()
                game.setup_empty_board()
                continue
            break

        if (boardw * boardh <= (piecesw+2) * (piecesh+2) or
            boardw < piecesw+2 or boardh < piecesh+2):
            pass    # not enough space
        elif bumpyborder:
            for x in range(1, game.bwidth-1):
                game.set(x, 1, '.')
                game.set(x, game.bheight-2, '.')
            for y in range(2, game.bheight-2):
                game.set(1, y, '.')
                game.set(game.bwidth-2, y, '.')

        boardname = 'CURRENT'
        boardinfo = {}
        initialboard = None
        if 'CURRENT' in getboarddict():
            del getboarddict()['CURRENT']   # to move it to the end
    else:
        boardname, boardinfo, initialboard = preset
        for (x, y), piece in initialboard.items():
            game.set(x, y, piece)
            if piece in player2bk:
                game.startpos = x, y

    px, py = game.startpos
    if game.board[px, py] not in player2bk:
        while game.board[px, py] not in bk2player:
            px += 1
        game.set(px, py, bk2player[game.board[px, py]])

    if initialboard is None:
        initialboard = game.board.copy()
        saveboard()

    undolist = []
    game.history = {}
    lastundolistlen = None
    starttime = time.time()

    while not game.finished(nb_pieces):
        game.flip()

        if CHEAT_SOLUTION_DB and lastundolistlen != len(undolist):
            prevcolor = game.COLOR['@']
            game.COLOR = game.__class__.COLOR.copy()
            if len(undolist) > 0:
                import db
                if not db.exists(game.board):
                    title = "Impossible"
                    game.COLOR['@'] = game.COLOR['&'] = (80, 0, 0)
            lastundolistlen = len(undolist)
            if prevcolor != game.COLOR['@']:
                game.set(px, py, game.board[px, py])
                game.flip()

        e = nextevent()
        if e.type == KEYDOWN:
            if e.key in KEY2DIR:
                dx, dy = KEY2DIR[e.key]
                pcar = game.board[px, py]
                game.set(px, py, player2bk[pcar])
                if game.board[px+dx, py+dy] not in bk2player:
                    blocked = [(px, py)]
                    bp = {}
                    while bp is not None and blocked:
                        for bx, by in blocked:
                            if not isinstance(game.board[bx+dx, by+dy], tuple):
                                bp = None
                                break
                            bp.update(game.bigpiece(bx+dx, by+dy))
                        else:
                            del blocked[:]
                            for x, y in bp:
                                if (x+dx, y+dy) not in bp:
                                    if game.board[x+dx, y+dy] != ' ':
                                        blocked.append((x, y))
                    if bp is not None:
                        undolist.append((px, py, pcar, game.history))
                        game.history = {}
                        for x, y in bp.keys():
                            bp[x, y] = game.board[x, y]
                            game.set(x, y, ' ')
                        for (x, y), piece in bp.items():
                            game.set(x+dx, y+dy, piece)
                if game.board[px+dx, py+dy] in bk2player:
                    px += dx
                    py += dy
                game.set(px, py, bk2player[game.board[px, py]])
            elif e.key == UNDO_KEY and undolist:
                bp = game.history.items()
                for (x, y), piece in bp:
                    game.set(x, y, ' ')
                for (x, y), piece in bp:
                    game.set(x, y, piece)
                px, py, pcar, game.history = undolist.pop()
                game.set(px, py, pcar)
        elif e.type == MOUSEBUTTONDOWN and CHEAT_MOUSE_CLICK:
            try:
                bp = game.bigpiece(*game.screen2xy(e.pos))
            except KeyError:
                bp = None
            if bp:
                pcar = game.board[px, py]
                undolist.append((px, py, pcar, game.history))
                game.history = {}
                bp = list(bp)
                random.shuffle(bp)
                for x, y in bp:
                    piece = game.board[x, y]
                    game.set(x, y, '#')
                    game.insert_piece(piece, pieces_max)
                    game.set(x, y, ' ')

    totalmoves = len(undolist)
    totaltime = int(time.time() - starttime)
    save = False
    if boardname.upper() == 'CURRENT':
        del getboarddict()['CURRENT']
        boardname = getboarddict().findfreename()
        save = True
    if not boardinfo.get('possible'):
        boardinfo['possible'] = True
        save = True
    if 'besttime' not in boardinfo or totaltime < boardinfo['besttime']:
        boardinfo['besttime'] = totaltime
        save = True
    if 'bestmoves' not in boardinfo or totalmoves < boardinfo['bestmoves']:
        boardinfo['bestmoves'] = totalmoves
        save = True
    if save:
        saveboard()
    ending(game, img, largeimg, game.finished(nb_pieces))


def ending(game, smallimg, largeimg, origin):
    # ending animation
    largeimg = largeimg.convert()
    starttime = time.time()
    pygame.display.flip()
    x1, y1 = game.xy2screen(origin)
    elapsed = 0.0
    smallimgw, smallimgh = smallimg.get_size()
    largeimgw, largeimgh = largeimg.get_size()
    factor0 = float(smallimgw+smallimgh)/(largeimgw+largeimgh)
    targetx, targety = game.screen.get_size()
    targetx = (targetx - largeimgw) // 2
    targety = (targety - largeimgh) // 2
    while elapsed < 1.0:
        elapsed = min(1.0, time.time() - starttime)
        factor = factor0 + (1.0-factor0) * elapsed
        nimage = pygame.transform.scale(largeimg, (int(factor*largeimgw),
                                                   int(factor*largeimgh)))
        for (x, y), value in game.board.items():
            if isinstance(value, tuple):
                value = ' '
            game.set(x, y, value)
        game.screen.blit(nimage, (int(elapsed*targetx + (1.0-elapsed)*x1),
                                  int(elapsed*targety + (1.0-elapsed)*y1)))
        pygame.display.flip()
    while True:
        nextevent()

# ____________________________________________________________

class BoardDict:

    def __init__(self, filename):
        self.filename = filename
        f = open(filename, 'r')
        lines = f.readlines()
        f.close()
        self.itemstext = ['']
        for i in range(len(lines)):
            line = lines[i].rstrip()
            if not line:
                if self.itemstext[-1]:
                    self.itemstext.append('')
            else:
                self.itemstext[-1] += line + '\n'
        if not self.itemstext[-1]:
            del self.itemstext[-1]
        self.name2index = {}
        for index in range(len(self)):
            self.name2index[self._getname(index)] = index

    def __len__(self):
        return len(self.itemstext)

    def keys(self):
        return [self._getname(index) for index in range(len(self))]

    def __contains__(self, name):
        return name.upper() in self.name2index

    def __getitem__(self, name):
        index = self.name2index[name.upper()]
        return self._decode(self.itemstext[index])

    def __setitem__(self, name, item):
        name = name.upper()
        text = self._encode(item, name)
        try:
            index = self.name2index[name]
        except KeyError:
            index = len(self.itemstext)
            self.itemstext.append(text)
            self.name2index[name] = index
        else:
            self.itemstext[index] = text
        self._flush()

    def __delitem__(self, name):
        index = self.name2index.pop(name.upper())
        del self.itemstext[index]
        self._flush()

    def findfreename(self):
        numbers = [-1]
        for key in self.name2index:
            if key.startswith('B') and key[1:].isdigit():
                info, board = self[key]
                numbers.append(int(key[1:]))
        name = 'B%d' % (max(numbers)+1,)
        return name

    def _flush(self):
        data = '\n'.join(self.itemstext)
        f = open(self.filename, 'w')
        f.write(data)
        f.close()

    def _getname(self, index):
        d = {}
        exec self.itemstext[index] in d
        names = [key for key in d if not key.startswith('_')]
        if len(names) == 1:
            return names[0].upper()
        else:
            return '?'

    def _decode(self, text):
        d = {}
        exec text in d
        [name] = [key for key in d if not key.startswith('_')]
        info, boardlist = d[name]
        assert isinstance(info, dict)
        assert isinstance(boardlist, list)
        board = {}
        for y in range(len(boardlist)):
            boardline = boardlist[y]
            for x in range(len(boardline)):
                board[x, y] = boardline[x]
        return info, board

    def _encode(self, (info, board), name):
        import pprint
        assert isinstance(info, dict)
        assert isinstance(board, dict)

        tiles = [tile for tile in board.values() if isinstance(tile, tuple)]
        wmax = max([w for w, h in tiles])
        hmax = max([h for w, h in tiles])
        info['tile'] = (wmax+1, hmax+1)

        bwidth, bheight = boardsize(board)
        cells = [repr(board.get((x, y), '*'))
                 for y in range(bheight) for x in range(bwidth)]
        w = max([len(s) for s in cells])
        lines = []
        it = iter(cells)
        for y in range(bheight):
            cellline = [it.next().center(w) for x in range(bwidth)]
            line = '  [%s],' % (', '.join(cellline),)
            lines.append(line)
        lines[0] = ' [' + lines[0][2:]
        lines[-1] = lines[-1][:-1] + '])'
        lines.insert(0, '%s = (%s,' % (name, pprint.pformat(info),))
        lines.append('')
        text = '\n'.join(lines)
        return text

def boardsize(board):
    keys = [key for key, value in board.items() if value != '*']
    xmax = max([x for x, y in keys])
    ymax = max([y for x, y in keys])
    return xmax+1, ymax+1

def getboarddict(FILENAME = 'boards.py', cache = {}):
    try:
        bdict = cache[FILENAME]
    except KeyError:
        bdict = cache[FILENAME] = BoardDict(FILENAME)
    return bdict

# ____________________________________________________________

def loadfile(filename, rescaled=None):
    # try to use the PIL to load the image file
    try:
        import PIL.Image
    except ImportError:
        pass
    else:
        image = PIL.Image.open(filename)
        try:
            image.load()
        except (IOError, IndexError):
            pass
        else:
            if image.mode != "RGB":
                image = image.convert("RGB")
            if rescaled:
                rw, rh = rescaled
                f = min(float(rw) / image.size[0],
                        float(rh) / image.size[1])
                rw = f * image.size[0]
                rh = f * image.size[1]
                image = image.resize((int(rw), int(rh)), PIL.Image.ANTIALIAS)
            size = image.size
            data = image.tostring()
            return pygame.image.frombuffer(data, size, "RGB")
    #
    if rescaled:
        # try to use 'convert'
        import os, tempfile
        tmpname = tempfile.mktemp('.bmp')
        try:
            rw, rh = rescaled
            os.system("convert '%s' -geometry %dx%d '%s'" % (
                filename, int(rw), int(rh), tmpname))
            try:
                f = open(tmpname, 'rb')
            except (IOError, OSError):
                pass
            else:
                try:
                    return pygame.image.load(f)
                finally:
                    f.close()
        finally:
            try:
                os.unlink(tmpname)
            except OSError:
                pass
    f = open(filename, 'rb')
    try:
        s = pygame.image.load(f)
    finally:
        f.close()
    if rescaled:
        rw, rh = rescaled
        sw, sh = s.get_size()
        ratio = min(float(rw)/sw, float(rh)/sh)
        rw = int(sw * ratio)
        rh = int(sh * ratio)
        s = pygame.transform.scale(s, (rw, rh))
    return s

def parse_args(argv):
    import sys, getopt
    opts, args = getopt.getopt(argv, "bh", ["help"])
    bumpyborder = ('-b', '') in opts
    if bumpyborder:
        opts.remove(('-b', ''))

    if not (1 <= len(args) <= 3):
        print >> sys.stderr, __doc__
        sys.exit(2)

    filename = args[0]
    if len(args) == 2 and not args[1][0].isdigit():
        bname = args[1]
        bdict = getboarddict()
        if bname not in bdict:
            print >> sys.stderr, "no board named %r in %s" % (
                bname, bdict.filename)
            sys.exit(1)
        info, board = bdict[bname]
        pw, ph = info['tile']
        bw, bh = boardsize(board)
        bw -= 2
        bh -= 2
        print '%s: %dx%d %dx%d' % (bname, pw, ph, bw, bh)
        if 'comment' in info:
            print info['comment']
        preset = bname, info, board
        if bumpyborder:
            print >> sys.stderr, "note: -b option ignored"
    else:
        if len(args) <= 1:
            pw, ph = 3, 3
        else:
            pw, ph = args[1].split('x')
            pw, ph = int(pw), int(ph)
        if len(args) <= 2:
            bw, bh = pw+3, ph+3
        else:
            bw, bh = args[2].split('x')
            bw, bh = int(bw), int(bh)
        preset = None

    largeimg = loadfile(filename)
    if FULLSCREEN_RES:
        tw, th = FULLSCREEN_RES
    else:
        tw, th = largeimg.get_size()
        if (bw+0.4)/pw > (bh+0.4)/ph:
            tw = tw * 5 // 4
        else:
            th = th * 5 // 4
    rw = tw/(bw+0.4)*pw
    rh = th/(bh+0.4)*ph
    smallimg = loadfile(filename, (rw, rh))
    if FULLSCREEN_RES:
        f = min(float(tw) / largeimg.get_size()[0],
                float(th) / largeimg.get_size()[1])
        if f < 0.99:
            largeimg = loadfile(filename, (largeimg.get_size()[0] * f,
                                           largeimg.get_size()[1] * f))
    return (pw, ph), (bw, bh), smallimg, largeimg, bumpyborder, preset

if __name__ == '__main__':
    import sys
    puzzle(*parse_args(sys.argv[1:]))