#!/usr/bin/env python3
# Assignment n 2, given in lesson 05 on 15/10/19
# Author: Leonardo Tamiano

import math
import numpy as np

# import previous work
from assignment_1.entropy import entropy

# Write a function called join_entropy() which computes the joint
# entropy of two random variables given their distribution mass
# functions.
#
# ARGUMENTS:
#
# joint_pmf is a numpy matrix describing the joint probabily mass
# functions of X and Y. In particular
# 
#            joint_pmf[i][j] = P(X = x_i, Y = y_j)
# 
def joint_entropy(joint_pmf):
    N_x, N_y = joint_pmf.shape
    result = 0
    
    for i in range(0, N_x):
        for j in range(0, N_y):
            if joint_pmf[i][j] != 0:
                result += joint_pmf[i][j] * math.log(joint_pmf[i][j], 2)
    return -result

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

# Write a function called conditional_entropy() that computes the
# conditional entropy function of variables X and Y, H(X|Y), given
# their joint p.m.f. and their marginal p.m.f.
#
# ARGUMENTS:
#
# 1) joint_pmf is a numpy matrix describing the joint probabily mass
#    functions of X and Y. In particular,
# 
#            joint_pmf[i][j] = P(X = x_i, Y = y_j)
# 
# 2) Y_pmf is a numpy vector describing the probabily mass function of
#    Y. In particular,
#
#            Y_pmf[j] = P(Y = y_j)
#
def conditional_entropy(joint_pmf, Y_pmf):
    N_x, N_y = joint_pmf.shape
    result = 0

    for i in range(0, N_x):
        for j in range(0, N_y):
            # Compute the conditional probability of X given Y using
            # the formula
            # 
            # P(X = x_i | Y = Y_i) = P(X = x_i, Y = y_j)/P(Y = y_j)
            # 
            cond_prob = joint_pmf[i][j] / Y_pmf[j]
            if cond_prob != 0:
                result += joint_pmf[i][j] * math.log(cond_prob, 2)
                
    return -result

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

# Write a function called mutual_information() that computes the
# mutual information of two discrete random variables X and Y, given
# their joint and marginal p.m.f.
#
# ARGUMENTS:
# 
# 1) joint_pmf is a numpy matrix describing the joint probabily mass
#    functions of X and Y. In particular,
# 
#            joint_pmf[i][j] = P(X = x_i, Y = y_j)
#
# 2) X_pmf is a numpy vector describing the probabily mass function of
#    X. In particular,
#
#            X_pmf[i] = P(X = x_i)
#
# 3) Y_pmf is a numpy vector describing the probabily mass function of
#    Y. In particular,
#
#            Y_pmf[j] = P(Y = y_j)
#
def mutual_information(joint_pmf, X_pmf, Y_pmf):
    N_x, N_y = joint_pmf.shape
    result = 0

    for i in range(0, N_x):
        for j in range(0, N_y):
            arg = joint_pmf[i][j] / (X_pmf[i] * Y_pmf[j])
            if arg != 0:
                result += joint_pmf[i][j] * math.log(arg, 2)
                
    return result

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

# Compute normalized versions of conditional entropy, joint entropy
# and mutual information for the discrete case

def normalized_conditional_entropy(joint_pmf, X_pmf, Y_pmf):
    return conditional_entropy(joint_pmf, Y_pmf) / entropy(X_pmf)

def normalzied_joint_entropy(joint_pmf, X_pmf, Y_pmf):
    return 1 - (mutual_information(joint_pmf, X_pmf, Y_pmf) / (entropy(X_pmf) + entropy(Y_pmf)))

def normalized_mutual_information_type1(joint_pmf, X_pmf, Y_pmf):
    return mutual_information(joint_pmf, X_pmf, Y_pmf) / joint_entropy(joint_pmf)

def normalized_mutual_information_type2(joint_pmf, X_pmf, Y_pmf):
    return (entropy(X_pmf) + entropy(Y_pmf))/joint_entropy(joint_pmf)

def normalized_mutual_information_type3(joint_pmf, X_pmf, Y_pmf):
    return mutual_information(joint_pmf, X_pmf, Y_pmf)/math.sqrt(entropy(X_pmf) * entropy(Y_pmf))

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

# EXTRA: compute other measures discusses during lecture.

def entropy_correlation_coefficient(joint_pmf, X_pmf, Y_pmf):
    return math.sqrt(2*mutual_information(joint_pmf, X_pmf, Y_pmf) / (entropy(X_pmf) + entropy(Y_pmf)))

def symmetric_uncertainty(joint_pmf, X_pmf, Y_pmf):
    return 2*mutual_information(joint_pmf, X_pmf, Y_pmf)/(entropy(X_pmf) + entropy(Y_pmf))

def variation_of_information(joint_pmf, X_pmf, Y_pmf):
    return joint_entropy(joint_pmf) - mutual_information(joint_pmf, X_pmf, Y_pmf)

def normalized_variation_of_information(joint_pmf, X_pmf, Y_pmf):
    return variation_of_information(joint_pmf, X_pmf, Y_pmf) / joint_entropy(joint_pmf)