import numpy as np
from scipy.stats import dirichlet
import ternary

scale = 60

x = np.load("votes_distrib.npz")
alphas = x["alphas"].mean(axis=0)

def prob(p):
    p = np.array(p)
    print(p)
    if np.all(p>0):
        return dirichlet.logpdf(p, alphas)
    else:
        return 0

figure, tax = ternary.figure(scale=scale)
tax.heatmapf(prob, boundary=True, style="triangular")
tax.boundary(linewidth=2.0)
tax.set_title("Shannon Entropy Heatmap")

tax.savefig("votes_distrib.eps")