#!/usr/bin/python3

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

mariosolver.setupSpaces()
DPI = 96
plt.figure(figsize=(800/DPI, 600/DPI), dpi=DPI)
plt.xlim(94, 0)
plt.ylim(0, 1)
tempFile = open('temp2', 'w')
spaces = range(95)
coins = {0: [], 5: [], 10: [], 15: [], 20: []}
#spaces.reverse()
for space in spaces:
    (n, d) = mariosolver.getProbabilityOfSuccess(space, 0, (True, True), frozenset())
    (n1, d1) = mariosolver.getProbabilityOfSuccess(space, 5, (True, True), frozenset())
    (n2, d2) = mariosolver.getProbabilityOfSuccess(space, 10, (True, True), frozenset())
    (n3, d3) = mariosolver.getProbabilityOfSuccess(space, 15, (True, True), frozenset())
    (n4, d4) = mariosolver.getProbabilityOfSuccess(space, 20, (True, True), frozenset())
    coins[0].append(n/d)
    coins[5].append(n1/d1)
    coins[10].append(n2/d2)
    coins[15].append(n3/d3)
    coins[20].append(n4/d4)
    tempFile.write('%d %.10f %.10f %.10f %.10f %.10f\n' % (space, n/d, n1/d1, n2/d2, n3/d3, n4/d4))
tempFile.close()
plt.ylabel("Probability of winning this turn")
plt.xlabel("Starting square (squares away from princess)")
plt.yticks(np.linspace(0, 1, num=6, endpoint=True))
plt.axvline(22, color='black')
plt.axvline(46, color='black')
plt.axvline(70, color='black')
plt.axvline(1, linestyle='--', color='black', alpha=0.5)
plt.axvline(6, linestyle='--', color='black', alpha=0.5)
plt.axvline(12, linestyle='--', color='black', alpha=0.5)
plt.axvline(15, linestyle='--', color='black', alpha=0.5)
plt.axvline(20, linestyle='--', color='black', alpha=0.5)
plt.axvline(56, linestyle='--', color='black', alpha=0.5)
plt.axvline(58, linestyle='--', color='black', alpha=0.5)
plt.axvline(66, linestyle='--', color='black', alpha=0.5)
plt.text(85, 0.95, "World 1")
plt.text(61, 0.95, "World 2")
plt.text(37, 0.95, "World 3")
plt.text(13, 0.95, "World 4")
plt.plot(spaces, coins[0], label="0 coins")
plt.plot(spaces, coins[5], label="5 coins")
plt.plot(spaces, coins[10], label="10 coins")
plt.plot(spaces, coins[15], label="15 coins")
plt.plot(spaces, coins[20], label="20 coins")
plt.legend(loc='center left')
plt.savefig('plotspace.png', dpi=DPI)