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

import math
import numpy as np
import matplotlib.pyplot as pyplot

# import previous work
from assignment_1.entropy import entropy


# DESCRIPTION
# 
# This function computes the binary entropy for a given discrete random
# variable X with pmf given by (p0, 1 - p0). 
#
# ARGUMENTS:
#
# 1) p0 is a real value between 0 and 1 representing the probabilty
#    that X assumes the first value. That is,
#
#                        P(X = x_0) = p_0
#
def binary_entropy(p0):
    return entropy(np.array([p0, 1 - p0]))

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

# DESCRIPTION
#
# The following code is used to plot the binary entropy function as a
# function of p0.
#
t = np.linspace(0, 1, 20)
y = [binary_entropy(i) for i in t]

pyplot.plot(t, y)
pyplot.xticks(np.arange(0, 2, 0.1))
pyplot.axis([0, 1, 0, 1])
pyplot.axvline(x=0.5, color='r', linestyle='--')
pyplot.show()