import random
import itertools
import time, math, copy

cache = {}
MIN_CAMELSTEP = 1
MAX_CAMELSTEP = 3

def getAdder(length):
    return itertools.product(range(MIN_CAMELSTEP, MAX_CAMELSTEP+1), repeat=length)

def cycleThroughStates(camelState, adder):
    l = len(camelState)
    for a in adder:
        newState = list(camelState)
        for i in range(l):
            newState[i] += a[i]
        yield newState

# takes in a camel state (list of camel positions) and returns
# the expected value of how many moves are left
def calculateRaceDynamicImpl(camelState, adder, numSquares):
    if (max(camelState) >= numSquares):
        return 0
    camelState.sort()
    if (tuple(camelState) in cache):
        return cache[tuple(camelState)]
    bigTotal = 0
    count = 0
    for newState in cycleThroughStates(camelState, adder):
        bigTotal += calculateRaceDynamicImpl(newState, adder, numSquares)
        count += 1
    answer = bigTotal/float(count) + 1.0
    cache[tuple(camelState)] = answer
    return answer

def calculateRaceDynamic(numCamels, numSquares):
    adder = list(getAdder(numCamels))
    cache = {}
    return calculateRaceDynamicImpl([-1] * numCamels, adder, numSquares)

def simulateRace(numCamels, numSquares):
    totalSteps = 0
    ITERATIONS = 20000
    for i in range(ITERATIONS):
        camelPositions = [-1] * numCamels
        while (max(camelPositions) < numSquares):
            for j in range(numCamels):
                camelPositions[j] += random.randint(MIN_CAMELSTEP, MAX_CAMELSTEP)
            totalSteps += 1
    return float(totalSteps) / ITERATIONS

# Assumes there's enough room in boardPositions already
def simulateLegWithStacking(camelPositions, boardPositions):
    camelOrder = list(range(len(camelPositions)))
    random.shuffle(camelOrder)
    # first camel in list is on the bottom
    for camelNum in camelOrder:
        curPosition = camelPositions[camelNum]
        newPosition = curPosition + random.randint(MIN_CAMELSTEP, MAX_CAMELSTEP)
        positionInStack = boardPositions[curPosition].index(camelNum)
        camelsToMove = boardPositions[curPosition][positionInStack:]
        boardPositions[curPosition] = boardPositions[curPosition][:positionInStack]
        assert len(camelsToMove) >= 1
        for c in camelsToMove:
            camelPositions[c] = newPosition
            boardPositions[newPosition].append(c)

def findSimulatedWinnerOfLeg(initialPositions):
    boardPositions = list(initialPositions)
    numCamels = max([max(l, default=-1) for l in initialPositions], default=-1) + 1
    assert numCamels >= 2
    for i in range(numCamels*MAX_CAMELSTEP):
        boardPositions.append([])
    camelPositions = [0] * numCamels
    for (i, l) in enumerate(initialPositions):
        for camel in l:
            camelPositions[camel] = i
    assert min(camelPositions) == 0
    assert max(camelPositions) == len(initialPositions) - 1
    
    ITERATIONS = 200000
    firstPlaces = [0] * numCamels
    secondPlaces = [0] * numCamels
    for i in range(ITERATIONS):
        camelPositionsTemp = list(camelPositions)
        boardPositionsTemp = copy.deepcopy(boardPositions)
        simulateLegWithStacking(camelPositionsTemp, boardPositionsTemp)
        finishers = []
        topSpot = max(camelPositionsTemp)
        while len(finishers) < 2:
            if len(boardPositionsTemp[topSpot]) > 0:
                for x in boardPositionsTemp[topSpot][::-1]:
                    finishers.append(x)
            topSpot -= 1
        finishers = finishers[:2]
        firstPlaces[finishers[0]] += 1
        secondPlaces[finishers[1]] += 1
    for i in range(numCamels):
        firstPlaces[i] /= float(ITERATIONS)
        secondPlaces[i] /= float(ITERATIONS)
    return (firstPlaces, secondPlaces)

def simulateRaceWithStacking(numCamels, numSquares, returnDistribution=False):
    totalSteps = 0
    ITERATIONS = 20000
    legsDist = {}
    for i in range(ITERATIONS):
        # extra slack for the last leg
        boardPositions = [[] for i in range(numSquares + MAX_CAMELSTEP * numCamels)]
        camelPositions = [-1] * numCamels
        # first initialize, they don't move with stacking on the first turn
        steps = 1
        for i in range(numCamels):
            camelPositions[i] += random.randint(MIN_CAMELSTEP, MAX_CAMELSTEP)
            boardPositions[camelPositions[i]].append(i)
        while (max(camelPositions) < numSquares):
            simulateLegWithStacking(camelPositions, boardPositions)
            steps += 1
        totalSteps += steps
        if returnDistribution:
            if steps not in legsDist:
                legsDist[steps] = 0
            legsDist[steps] += 1
    if returnDistribution:
        listOfLegs = []
        for s in sorted(legsDist.keys()):
            listOfLegs.append([s-1, legsDist[s]/ITERATIONS])
        return listOfLegs
    return float(totalSteps) / ITERATIONS

def probFinishLessThanEqualDieRollsIter(curState, cached):
    if curState in cached:
        return cached[curState]
    numSquares = curState[0]
    dieRollsLeft = curState[1]
    if (numSquares <= 0):
        return 1.0
    if (dieRollsLeft <= 0):
        return 0.0
    totalProb = 0.0
    for step in range(MIN_CAMELSTEP, MAX_CAMELSTEP+1):
        totalProb += probFinishLessThanEqualDieRollsIter((numSquares - step, dieRollsLeft - 1), cached)
    totalProb *= 1.0 / (MAX_CAMELSTEP + 1 - MIN_CAMELSTEP)
    cached[curState] = totalProb
    return totalProb

# Calculates the probability that one camel will finish in
# less than or equal to this many die rolls.
def probFinishLessThanEqualDieRolls(numSquares, dieRolls, cached):
    state = (numSquares, dieRolls)
    return probFinishLessThanEqualDieRollsIter(state, cached)

def calculateRaceByDieRolls(numCamels, numSquares, returnDistribution=False):
    # the camels effectively start at square -1, so they actually have to go one extra square
    numSquares = numSquares + 1
    # figure out the range of die rolls that are possible
    minDieRolls = math.ceil(numSquares/MAX_CAMELSTEP)
    maxDieRolls = math.ceil(numSquares/MIN_CAMELSTEP)
    expected = 0.0
    lastProb = 0.0
    cached = {}
    listOfProbs = []
    for d in range(minDieRolls, maxDieRolls+1):
        prob = probFinishLessThanEqualDieRolls(numSquares, d, cached)
        probLEQThisWithAnyCamel = 1.0 - math.pow(1.0-prob, numCamels)
        assert probLEQThisWithAnyCamel <= 1.0
        probExactlyThisManyWithAnyCamel = probLEQThisWithAnyCamel - lastProb
        lastProb = probLEQThisWithAnyCamel
        if returnDistribution:
            listOfProbs.append([d-1, probExactlyThisManyWithAnyCamel])
        else:
            expected += d * probExactlyThisManyWithAnyCamel
    if returnDistribution:
        return listOfProbs
    else:
        return expected

def printWithTime(label, func):
    print(label)
    t1 = time.perf_counter()
    try:
        print("%.2f" % func())
    except Exception as e:
        print("CAUGHT EXCEPTION: %s" % str(e))
    t2 = time.perf_counter()
    print("time taken: %.2f secs" % (t2 - t1))

if (__name__ == '__main__'):
    NUM_CAMELS = 5
    NUM_SQUARES = 16

    #print(findSimulatedWinnerOfLeg([[0,1]]))
    #print(findSimulatedWinnerOfLeg([[0],[1]]))
    #print(findSimulatedWinnerOfLeg([[0,1],[2]]))
    #print(findSimulatedWinnerOfLeg([[0,1],[],[2]]))
    #print(findSimulatedWinnerOfLeg([[0,1,2],[3]]))
    #print(findSimulatedWinnerOfLeg([[0,1,2],[],[3]]))
    #print(findSimulatedWinnerOfLeg([[0,1],[2],[3]]))
    #print(findSimulatedWinnerOfLeg([[0,1],[2,3],[4]]))
    #print(findSimulatedWinnerOfLeg([[0,1,2],[3],[4]]))
    
    #printWithTime("simulateRaceWithStacking:", lambda: simulateRaceWithStacking(NUM_CAMELS, NUM_SQUARES) - 1)
    #printWithTime("simulateRace:", lambda: simulateRace(NUM_CAMELS, NUM_SQUARES) - 1)
    #printWithTime("calculateRaceDynamic:", lambda: calculateRaceDynamic(NUM_CAMELS, NUM_SQUARES) - 1)
    #printWithTime("calculateRaceByDieRolls:", lambda: calculateRaceByDieRolls(NUM_CAMELS, NUM_SQUARES) - 1)

    #print(simulateRaceWithStacking(NUM_CAMELS, NUM_SQUARES, True))
    #print(calculateRaceByDieRolls(NUM_CAMELS, NUM_SQUARES, True))
    #numCamelsToExpected = []
    #for c in range(1,101):
    #    numCamelsToExpected.append([c, calculateRaceByDieRolls(c, NUM_SQUARES) - 1])
    #print(numCamelsToExpected)

    #numCamelsToExpected = []
    #for c in range(1,101):
    #    numCamelsToExpected.append([c, simulateRaceWithStacking(c, NUM_SQUARES) - 1])
    #print(numCamelsToExpected)
