# CDI A310
# Pablo Llueca

# Load numpy and pyplot
import numpy as np
import matplotlib.pyplot as plt
#file_path = ""  # set your path for png images
file_path = "d:/Docencia/2015/labs/Llueca/"

#----------------------------------
# Functions from L310
#----------------------------------
# hline draws a horizontal line fron (a,h) to (b,h)
# The default width and dashing style are 1 and 'k-'
def hline(a,b,h,lw=1,dash='k-'):
    x = np.arange(a,b+0.0001,b-a)
    return plt.plot(x,h+0*x,dash, lw=lw)

# vline draws a vertical line fron (d,a) to (d,b)
# The default width and dashing style are 1 and 'k-'
def vline(a,b,d,lw=1, dash='k-'):
    y = np.arange(a,b+0.0001,b-a)
    return plt.plot(d+0*y,y,dash, lw=lw)

# To tick and label the vertical axis
def htick(X, T, run=0.03, lw='1'):
    plt.plot([X[0]-run, X[0]+run], [X[1],X[1]], lw=lw, color='black')
    write((X[0]-2*run,X[1]), T, ha='right', fs='13')
    return None

# To tick and label the horizontal axis
def vtick(X, T, run=0.03, lw='1'):
    plt.plot([X[0], X[0]], [X[1]-run,X[1]+run], lw=lw, color='black')
    write((X[0],X[1]-3*run), T, fs='13')
    return None

# Utility to write text T at position X with
# four parameters with default values
def write(X,T,ha='center', va = 'center', fs='16', rot=0.0):
    return plt.text(X[0],X[1], T, horizontalalignment=ha,verticalalignment=va, fontsize=fs, rotation=rot)



# function of first plot
def h(p): 
    l = []
    for x in p:
        if x == 0.0:
            l.append(0.0)
        else: l.append(-x*np.log2(x))
    return np.array(l)


# function of 2nd plot
def C(p): 
    l = []
    for x in p:
        if x==0.0 or x==1.0:
            l.append(1.0)
        else:
            l.append(1 + (x * np.log2(x)) + ((1-x)*np.log2(1-x)))
    return np.array(l)

# Graph 1: T2.4
def draw_graph1():
    fig1 = plt.figure("Entropy of a random experiment with probabilities (x, 1-x)")

    # graph of h_
    x = np.arange(0,1.0,0.001)
    plt.plot(x, h(x),'r-', lw=2)

    xmin = 0.0
    xmax = 1.0
    ymin = 0.0
    ymax = 0.6
    plt.xlim(xmin,xmax)
    plt.ylim(ymin,ymax)

    p=0.368; q=-p*np.log2(p)
    plt.plot([p],[q],'bo')
    hline(0,p,q,dash='g:')  
    vline(0,q,p,dash='g:')  


    write((p,-0.06), '$1/e \simeq 0.367$')  # 1/e
    write((0.9, 0.5),'$h(x)=-p\,\log_2(x)$', ha='right')
    write((p+0.01,q-0.4),'$\log_2(e)/e \simeq 0.53$', ha='left')
    
    plt.savefig(file_path + "T2_4.png", bbox_inches='tight')

def draw_graph2():
    fig2 = plt.figure("Capacity of a Binary Symmetric Chanel with Bitflip probability p")
    x = np.arange(0,1, 0.001)
    plt.plot(x, C(x),'r-', lw=2)  #function
    plt.plot([0,0.5, 1],[1,0, 1], 'go') 
    xmin = -0.20
    xmax = 1.1
    ymin = 0.0
    ymax = 1.1
    plt.xlim(xmin,xmax)
    plt.ylim(ymin,ymax)
    hline(0, 1.0, 1.0, dash='g:')
    vline(0, 1.0, 1.0, dash='g:')
    vline(0, 1.0, 0.0, dash='g:')   
    write((0.5, -0.1), '$C(p)= 1 + p\,\log_2(p)+ (1-p)\log_2(1-p)$')
    plt.savefig(file_path + "T2_17.png", bbox_inches='tight')


def main():
    plt.close("all")
    draw_graph1()
    draw_graph2()
    plt.show()

if __name__ == "__main__":
    main()