#!/usr/bin/env python3
# Assignment n 1, given in lesson 03 on 01/10/19
# Author: Leonardo Tamiano

import math
import random as r
import numpy as np
import matplotlib.pyplot as pyplot
from mpl_toolkits.mplot3d import Axes3D

# import previous work
from assignment_1.entropy import entropy

# DESCRIPTION
# 
# This function computes the ternary entropy for a given discrete random
# variable X with pmf given by (p0, p1, 1 - p0 - p1). 
#
# ARGUMENTS:
#
# 1) p0 is the probability that X assumes the first value. That is,
#
#                        P(X = x_0) = p_0
# 
# 2) p1 is the probability that X assumes the second value. That is,
#
#                        P(X = x_1) = p_1
#
def ternary_entropy(p0, p1):
    return entropy(np.array([p0, p1, 1 - p0 - p1]))

# ------------------------------------------------------------------

# DESCRIPTION
#
# The following code is used to plot the ternary entropy function as a
# function of p0 and p1.
#

STEPS = 20
couples = []

# construct a set of couples (p0, p1) such that p0 + p1 <= 1.
for p0 in np.linspace(0, 1, STEPS):
    if p0 == 1:
        # if p0 = 1, then p1 can only be 0.
        couples.append((1, 0))
    else:
        # otherwise p1 can range from 0 to 1 - p0
        for p1 in np.linspace(0, 1 - p0, STEPS):
            couples.append((p0, p1))
            
# construct triplets (X, Y, Z), such that Z = ternary_entropy(X, Y).
X = np.zeros(len(couples))
Y = np.zeros(len(couples))
Z = np.zeros(len(couples))

for i in range(0, len(couples)):
    x, y = couples[i]
    X[i] = (x)
    Y[i] = (y)
    Z[i] = (ternary_entropy(x, y))
    
# plot the computed data    
fig = pyplot.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_trisurf(X, Y, Z)
pyplot.show()