# Source: content/notes/math/information-theory.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np

def entropy(p):
    p = np.asarray(p, dtype=float)
    if np.any(p < 0) or not np.isclose(p.sum(), 1):
        raise ValueError("Expected normalized nonnegative masses")
    positive = p > 0
    return -np.sum(p[positive]*np.log2(p[positive]))

def kl(p, q):
    p, q = np.asarray(p, float), np.asarray(q, float)
    if p.shape != q.shape:
        raise ValueError("Shapes must agree")
    entropy(p)
    entropy(q)
    active = p > 0
    if np.any(q[active] == 0):
        return np.inf
    return np.sum(p[active]*np.log2(p[active]/q[active]))

joint = np.array([[3/8, 1/8], [1/8, 3/8]])
mi = kl(joint, joint.sum(1)[:, None]*joint.sum(0)[None, :])
assert np.isclose(mi, 1-entropy([.25, .75]))
later = 1-entropy([.375, .625])
assert 0 <= later < mi
assert entropy([1., 0.]) == 0.
assert np.isinf(kl([1., 0.], [0., 1.]))
p, q = np.array([.8, .2]), np.array([.5, .5])
assert not np.isclose(kl(p, q), kl(q, p))
lengths = np.array([1, 2, 3, 3])
probabilities = np.array([.5, .25, .125, .125])
assert np.isclose(np.sum(2.**-lengths), 1.)
assert np.isclose(probabilities@lengths, entropy(probabilities))
print("MI before / after second channel:", mi, later)