#!/usr/bin/env python
# -*- coding: iso-8859-15 -*-

from numpy import *
import sys
import os
import random
import pprint


outcomes = []

# roll 4d6
for d1 in range(1, 7):
    for d2 in range(1, 7):
        for d3 in range(1, 7):
            for d4 in range(1, 7):
outcome = array([d1, d2, d3, d4])
outcomes.append(outcome)

outcomes = array(outcomes)
hist = {}

# drop the lowest, count the number of times a given sum is found
for x in range(len(outcomes)):
    outcomes[x][outcomes[x].argmin()] = 0
    if outcomes[x].sum() in hist:
        hist[outcomes[x].sum()] += 1
    else:
        hist[outcomes[x].sum()] = 1

rsum = outcomes.sum()
sums = outcomes.sum(axis=1)
mean = sums.mean()
sigma = sums.std()

counts = array(hist.values())
hist2 = hist.copy()

# determine the probability of each sum occurring
for key,value in hist2.items():
    hist2[key] = counts[key-3] / float(counts.sum())

f = open("dice.data", "w")
for key in hist2.keys():
print >>f, "%d\t%0.12f" % (key, hist2[key])
f.close()

f = open("dice2.data", "w")
for key,value in hist2.items():
    print >>f, "%d\t%f" % (key, value)
f.close()

pprint.pprint(hist.values())

args = {
    "mean": mean,
    "sigma": sigma,
    "-sigma": mean - sigma,
    "+sigma": mean + sigma,
    "++sigma": mean + (2 * sigma),
    "--sigma": mean - (2 * sigma),
    "+++sigma": mean + (3 * sigma),
    "---sigma": mean - (3 * sigma),
}
labels = {
    "+sigma_label": mean + sigma - 0.1,
    "-sigma_label": mean - sigma + 0.1,
    "++sigma_label": mean + (2*sigma) - 0.1,
    "--sigma_label": mean - (2*sigma) + 0.1,
    "+++sigma_label": mean + (3*sigma) - 0.1,
    "---sigma_label": mean - (3*sigma) + 0.1,
    "mean_label": mean - 0.1,
}
args.update(labels)
gnuplot = """
set terminal png enhanced size 800,600 font '/Library/Fonts/Arial.ttf' 8
set output "dice.png"
set arrow from %(mean)f,0 to %(mean)f,0.14 nohead lc 1
set label "? (%(mean)0.2f)" at %(mean_label)f,0.05 right tc ls 1
set arrow from %(-sigma)f,0 to %(-sigma)f,0.14 nohead lc 2
set label "?? (%(-sigma)0.2f)" at %(-sigma_label)f,0.05 tc ls 2
set arrow from %(+sigma)f,0 to %(+sigma)f,0.14 nohead lc 2
set label "+? (%(+sigma)0.2f)" at %(+sigma_label)f,0.05 right tc ls 2

set arrow from %(--sigma)f,0 to %(--sigma)f,0.14 nohead lc 2
set label "?2? (%(--sigma)0.2f)" at %(--sigma_label)f,0.05 tc ls 2
set arrow from %(++sigma)f,0 to %(++sigma)f,0.14 nohead lc 2
set label "+2? (%(++sigma)0.2f)" at %(++sigma_label)f,0.05 right tc ls 2

set arrow from %(---sigma)f,0 to %(---sigma)f,0.14 nohead lc 3
set label "?3? (%(---sigma)0.2f)" at %(---sigma_label)f,0.05 tc ls 3
set arrow from %(+++sigma)f,0 to %(+++sigma)f,0.14 nohead lc 3
set label "+3? (%(+++sigma)0.2f)" at %(+++sigma_label)f,0.05 right tc ls 3

plot 'dice.data' with lines smooth csplines title '4d6 drop lowest'
""" % args

f = open("dice.gnuplot", "w")
f.write(gnuplot)
f.close()

gnuplot = """
set terminal png enhanced size 800,600 font '/Library/Fonts/Arial.ttf' 8
set output "dice2.png"
set boxwidth 0.9 relative
set style data histograms
set style fill solid 1.0 border -1
set style data histogram
plot 'dice2.data' using 2:xticlabels(1) title '4d6 drop lowest'
"""

f = open("dice2.gnuplot", "w")
f.write(gnuplot)
f.close()

print "mean: %f" % mean
print "sigma: %f" % sigma
print "1 sigma: %0.4f - %0.4f" % (mean - sigma, mean + sigma)
print "2 sigma: %0.4f - %0.4f" % (mean - 2*sigma, mean + 2*sigma)
print "3 sigma: %0.4f - %0.4f" % (mean - 3*sigma, mean + 3*sigma)

os.system("gnuplot 'dice.gnuplot'")
os.system("gnuplot 'dice2.gnuplot'")