#!/usr/bin/env python

"""Wrapper script for testing the performance of simple AI systems.

bm_ai.py runs the following little solvers:
    - N-Queens

This used to contain an alphametics solver, but that was found to be bound
primarily by eval() performance.
"""

# borrowed from pypy/unladen_swallow repos

__author__ = "collinwinter@google.com (Collin Winter)"

# Python imports
import time

# Local imports
import util

# Pure-Python implementation of itertools.permutations().
def permutations(iterable, r=None):
    """permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)"""
    pool = tuple(iterable)
    n = len(pool)
    if r is None:
        r = n
    indices = list(range(n))
    cycles = list(range(n-r+1, n+1))[::-1]
    yield tuple(pool[i] for i in indices[:r])
    while n:
        for i in reversed(range(r)):
            cycles[i] -= 1
            if cycles[i] == 0:
                indices[i:] = indices[i+1:] + indices[i:i+1]
                cycles[i] = n - i
            else:
                j = cycles[i]
                #fix me brython (this causes issues), but I cannot
                #replicate them in the editor (maybe its a generator issue)?
                #indices[i], indices[-j] = indices[-j], indices[i]
                _temp=indices[i]
                indices[i]=indices[-j]
                indices[j]=_temp
                yield tuple(pool[i] for i in indices[:r])
                break
        else:
            return


# From http://code.activestate.com/recipes/576647/
def n_queens(queen_count):
    """N-Queens solver.

    Args:
        queen_count: the number of queens to solve for. This is also the
            board size.

    Yields:
        Solutions to the problem. Each yielded value is looks like
        (3, 8, 2, 1, 4, ..., 6) where each number is the column position for the
        queen, and the index into the tuple indicates the row.
    """

    cols = list(range(queen_count))
    for vec in permutations(cols):
        if (queen_count == len(set(vec[i]+i for i in cols))
                        == len(set(vec[i]-i for i in cols))):
            yield vec


def test_n_queens(iterations):
    from javascript import console
    # Warm-up runs.
    list(n_queens(8))
    list(n_queens(8))

    times = []
    for _ in range(iterations):
        t0 = time.time()
        list(n_queens(8))
        t1 = time.time()
        times.append(t1 - t0)
    return times

def run(num_runs=100, take_geo_mean=True):
    return util.run_benchmark(take_geo_mean, num_runs, test_n_queens)
