#!/usr/bin/python3

import mariosolver
import matplotlib
# Turn off DISPLAY
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

mariosolver.setupSpaces()
plt.xlim(0, 100)
plt.ylim(0, 1)
tempFile = open('temp', 'w')
normal = []
noPickUpCoins = []
noObstacles = []
noObstaclesOrPickUpCoins = []
for c in range(101):
    (n, d) = mariosolver.getProbabilityOfSuccess(94, c, (True, True), frozenset())
    (n1, d1) = mariosolver.getProbabilityOfSuccess(94, c, (True, False), frozenset())
    (n2, d2) = mariosolver.getProbabilityOfSuccess(94, c, (False, True), frozenset())
    (n3, d3) = mariosolver.getProbabilityOfSuccess(94, c, (False, False), frozenset())
    tempFile.write('%d %.10f %.10f %.10f %.10f\n' % (c, n/d, n1/d1, n2/d2, n3/d3))
    normal.append(n/d)
    noPickUpCoins.append(n1/d1)
    noObstacles.append(n2/d2)
    noObstaclesOrPickUpCoins.append(n3/d3)
tempFile.close()
plt.ylabel("Probability of winning")
plt.xlabel("Starting number of coins")
plt.yticks(np.linspace(0, 1, num=11, endpoint=True))
plt.xticks(np.linspace(0, 100, num=6, endpoint=True))
plt.grid(axis='both', linestyle='-', color='0.5')
plt.plot(range(101), normal, label="Normal")
plt.plot(range(101), noPickUpCoins, label="No pick up coins")
plt.plot(range(101), noObstacles, label="No obstacles")
plt.plot(range(101), noObstaclesOrPickUpCoins, label="No obstacles/pick up coins")
plt.legend(loc='lower right')
plt.savefig('plot.png')