# Written by Bram Cohen
# see LICENSE.txt for license information
#
# $Id: PiecePicker.py 68 2006-04-26 20:14:35Z sgrayban $
#

from random import randrange, shuffle, choice

class PiecePicker:
    def __init__(self, numpieces, rarest_first_cutoff = 1):
        self.rarest_first_cutoff = rarest_first_cutoff
        self.numpieces = numpieces
        self.interests = [range(numpieces)]
        self.pos_in_interests = range(numpieces)
        self.numinterests = [0] * numpieces
        self.started = []
        self.seedstarted = []
        self.numgot = 0
        self.scrambled = range(numpieces)
        shuffle(self.scrambled)

    def got_have(self, piece):
        if self.numinterests[piece] is None:
            return
        numint = self.numinterests[piece]
        if numint == len(self.interests) - 1:
            self.interests.append([])
        self.numinterests[piece] += 1
        self._shift_over(piece, self.interests[numint], self.interests[numint + 1])

    def lost_have(self, piece):
        if self.numinterests[piece] is None:
            return
        numint = self.numinterests[piece]
        self.numinterests[piece] -= 1
        self._shift_over(piece, self.interests[numint], self.interests[numint - 1])

    def _shift_over(self, piece, l1, l2):
        p = self.pos_in_interests[piece]
        l1[p] = l1[-1]
        self.pos_in_interests[l1[-1]] = p
        del l1[-1]
        newp = randrange(len(l2) + 1)
        if newp == len(l2):
            self.pos_in_interests[piece] = len(l2)
            l2.append(piece)
        else:
            old = l2[newp]
            self.pos_in_interests[old] = len(l2)
            l2.append(old)
            l2[newp] = piece
            self.pos_in_interests[piece] = newp

    def requested(self, piece, seed = False):
        if piece not in self.started:
            self.started.append(piece)
        if seed and piece not in self.seedstarted:
            self.seedstarted.append(piece)

    def complete(self, piece):
        assert self.numinterests[piece] is not None
        self.numgot += 1
        l = self.interests[self.numinterests[piece]]
        p = self.pos_in_interests[piece]
        l[p] = l[-1]
        self.pos_in_interests[l[-1]] = p
        del l[-1]
        self.numinterests[piece] = None
        try:
            self.started.remove(piece)
            self.seedstarted.remove(piece)
        except ValueError:
            pass

    def next(self, havefunc, seed = False):
        bests = None
        bestnum = 2 ** 30
        if seed:
            s = self.seedstarted
        else:
            s = self.started
        for i in s:
            if havefunc(i):
                if self.numinterests[i] < bestnum:
                    bests = [i]
                    bestnum = self.numinterests[i]
                elif self.numinterests[i] == bestnum:
                    bests.append(i)
        if bests:
            return choice(bests)
        if self.numgot < self.rarest_first_cutoff:
            for i in self.scrambled:
                if havefunc(i):
                    return i
            return None
        for i in xrange(1, min(bestnum, len(self.interests))):
            for j in self.interests[i]:
                if havefunc(j):
                    return j
        return None

    def am_I_complete(self):
        return self.numgot == self.numpieces

    def bump(self, piece):
        l = self.interests[self.numinterests[piece]]
        pos = self.pos_in_interests[piece]
        del l[pos]
        l.append(piece)
        for i in range(pos,len(l)):
            self.pos_in_interests[l[i]] = i

def test_requested():
    p = PiecePicker(9)
    p.complete(8)
    p.got_have(0)
    p.got_have(2)
    p.got_have(4)
    p.got_have(6)
    p.requested(1)
    p.requested(1)
    p.requested(3)
    p.requested(0)
    p.requested(6)
    v = _pull(p)
    assert v[:2] == [1, 3] or v[:2] == [3, 1]
    assert v[2:4] == [0, 6] or v[2:4] == [6, 0]
    assert v[4:] == [2, 4] or v[4:] == [4, 2]

def test_change_interest():
    p = PiecePicker(9)
    p.complete(8)
    p.got_have(0)
    p.got_have(2)
    p.got_have(4)
    p.got_have(6)
    p.lost_have(2)
    p.lost_have(6)
    v = _pull(p)
    assert v == [0, 4] or v == [4, 0]

def test_change_interest2():
    p = PiecePicker(9)
    p.complete(8)
    p.got_have(0)
    p.got_have(2)
    p.got_have(4)
    p.got_have(6)
    p.lost_have(2)
    p.lost_have(6)
    v = _pull(p)
    assert v == [0, 4] or v == [4, 0]

def test_complete():
    p = PiecePicker(1)
    p.got_have(0)
    p.complete(0)
    assert _pull(p) == []
    p.got_have(0)
    p.lost_have(0)

def test_rarer_in_started_takes_priority():
    p = PiecePicker(3)
    p.complete(2)
    p.requested(0)
    p.requested(1)
    p.got_have(1)
    p.got_have(0)
    p.got_have(0)
    assert _pull(p) == [1, 0]

def test_zero():
    assert _pull(PiecePicker(0)) == []

def _pull(pp):
    r = []
    def want(p, r = r):
        return p not in r
    while True:
        n = pp.next(want)
        if n is None:
            break
        r.append(n)
    return r
