#!/usr/bin/env python
#-*- coding:utf-8 -*-
"""
AI-related functions
"""

# XXX take penalty into account, in a smart way:
# threshold based on number of remaining blocks ?
# threshold based on the number of variations that can be played
# for each block (a block which can only be placed in one way
# might to interesting to play immediately)

from numpy.core import hstack, vstack
from numpy.core.numeric import array, argsort, mean
from numpy.core.fromnumeric import argmax
from matplotlib.mlab import find
from itertools import imap, ifilter, izip

from blokus3d.utils import randomFromList, fst, snd, randint, unik

randomMove = lambda gs : randomFromList(gs.legalMoves())

# Fitness functions

def relativeBaseScoreFitness(gs):
    scores = gs.baseScores()
    currentPlayer = (gs.nextPlayer-1) % gs.nbPlayers
    # Remove player score from the scores
    bestOpponentScore = max(hstack([scores[:currentPlayer],\
                            scores[currentPlayer+1:]]))
    return scores[currentPlayer] - bestOpponentScore

def libertiesFitness(gs):
    """Each legal cube amounts to its square height"""
    return sum(cube[2]**2 for cube in \
               gs.libertyCubes([(gs.nextPlayer-1) % gs.nbPlayers]))

def penaltyFitness(gs):
    return -gs.penalty()[(gs.nextPlayer-1) % gs.nbPlayers]

def mixtureFitness(weightedFitFuns):
    """Meta-fitness function that make a mix between several"""
    return lambda gs: sum(weight*fitFun(gs) \
                           for (weight,fitFun) in weightedFitFuns)

def bestMoves(gs, fitFun, moves=None):
    """Select the best legal moves according to
       a one-step fitness function"""
    if moves == None:
        moves = gs.legalMoves()
    if moves == []:
        return [None], None
    fitnesses = array([fitFun(gs.clone().playMove(move)) \
                       for move in moves])
    bestFitness = max(fitnesses)
    # TODO should use argsort instead
    selectedMoves = map(moves.__getitem__,find(fitnesses==bestFitness))
    return selectedMoves, bestFitness

# Heuristics based on fitness functions

def oneStepHeuristic(gs, fitFuns, verbose=False):
    """Find the move that maximize the fitness
       after this move."""
    moves = gs.legalMoves()
    for fitFun in fitFuns:
        moves, fitness = bestMoves(gs, fitFun)
        if verbose:
            print "best %s : %d (%d moves)" \
                    % (str(fitFun), fitness, len(moves))
        if len(moves)==1:
            return moves[0]
    return moves[0]

mixtureOneStepHeuristic = lambda gs, weightedFitFuns:\
        oneStepHeuristic(gs, [mixtureFitness(weightedFitFuns)])

def threeHeuristicsMC(gs, verbose=False):
    """Find the best move according to three heuristics,
       then evaluate each in depth with Monte-Carlo method"""
    lm = gs.legalMoves()
    if lm == []:
        yield None
        raise StopIteration
    if verbose:
        print "Picking a best next move according to each criterion"
    moves = map(lambda fitFun: fst(fst(bestMoves(gs,fitFun,moves=lm))),\
                [relativeBaseScoreFitness,\
                 libertiesFitness,\
                 penaltyFitness])
    # Some moves might be the same
    moves = unik(moves)
    if verbose:
        print "Candidate moves are :",moves
    assert moves != []
    if len(moves)==1:
        yield moves[0]
        raise StopIteration
    if verbose:
        print "Now performing MC evaluation"
    for move in monteCarloHeuristic(gs,moves=moves):
        yield move

# Meta-heuristic

# TODO add depth so it can be maximini...maximinimax !
def minimax(gs, fitFun):
    # Get the legal moves and return immediately if there are only one
    moves = gs.legalMoves()
    lnMoves = len(moves)
    if lnMoves==1:
        yield moves[0]
        raise StopIteration
    # Get the one-step fitness for each possible move
    fitnesses = []
    nextGss = []
    for move in moves:
        nextGs = gs.clone().playMove(move)
        fitnesses.append(fitFun(nextGs))
        nextGss.append(nextGs)
    bestMovesOrder = argsort(fitnesses)[::-1]
    bestSoFar = moves[bestMovesOrder[0]]
    yield bestSoFar
    # Find the one-step that minimize the opponent fitness (second step)
    leastBestEnemyFitness = snd(bestMoves(nextGss[bestMovesOrder[0]].clone(),fitFun))
    for num,idx in enumerate(bestMovesOrder):
        print "processing move %d/%d (fitness %d)" % (num+1,lnMoves,fitnesses[idx])
        bestEnemyFitness = snd(bestMoves(nextGss[idx].clone(),fitFun))
        if bestEnemyFitness < leastBestEnemyFitness:
            bestSoFar = moves[idx]
            leastBestEnemyFitness = bestEnemyFitness
            print "new least best enemy fitness : %d" % leastBestEnemyFitness
        # TODO the yield is only here to avoid blocking on timeLimit,
        # must be removed when timeLimit works asynchronously
        yield bestSoFar

# Some other functions

def monteCarloScores(gs, maxDepth=None):
    # Copy the game state, so we can keep the original
    gs = gs.clone()
    depth = 0
    while not gs.isOver() and (maxDepth==None or depth<maxDepth):
        lm = gs.legalMoves()
        if lm != []:
            choice = randint(len(lm))
            gs.playMove(lm[choice])
            depth += 1
        else:
            gs.playMove(None) # passing move
    return gs.finalScores()

def monteCarloHeuristic(gs, moves=None, maxDepth=None, verbose=False):
    """Determines the best move using a Monte-Carlo estimation
       of the final scores"""
    if moves==None:
        moves = gs.legalMoves()
    nextStates = map(lambda move : gs.clone().playMove(move), moves)
    relativeScoreGrid = [[] for _ in xrange(len(moves))]
    for trial in xrange(1000000):
        for m in xrange(len(moves)):
            scores = nextStates[m].monteCarloScores(\
                        gs,maxDepth=maxDepth)
            if verbose:
                print "trial %d, move %d/%d, scores : %s" \
                      % (trial+1, m+1, len(moves), scores)
            relativeScoreGrid[m].append(scores)
        # for each move, compute the margin of each winning trial
        winningMargin = [ map(lambda z : z[1][-1]-z[1][-2], \
            ifilter(lambda y: y[0][-1]==gs.nextPlayer, \
            izip(imap(argsort,x),x))) \
            for x in relativeScoreGrid ]
        choice = argmax(map(mean,winningMargin))
        if verbose:
            print "best move so far is %d, mean margin = %f (wins %d/%d)" \