import numpy as np

# ---- linear autoencoder: 3 -> 1 -> 3 (matches the lecture's worked example) ----
x  = np.array([1.0, 0.5, -0.5])
We = np.array([0.4, 0.3, -0.2])          # encoder weights (3,)
be = 0.1
Wd = np.array([0.5, 0.4, -0.3])          # decoder weights (3,)
bd = np.array([0.05, 0.02, -0.01])

def encode(inp): return inp @ We + be
def decode(z):    return z * Wd + bd
def mse(a, b):     return np.mean((a - b) ** 2)

# ---- clean forward pass ----
z     = encode(x)
x_hat = decode(z)
loss  = mse(x, x_hat)
print("Clean:    z =", z, " x_hat =", x_hat, " MSE =", round(loss, 5))

# ---- denoising: masking noise zeroes feature x2 (index 1) ----
x_tilde     = x.copy(); x_tilde[1] = 0.0
z_tilde     = encode(x_tilde)
x_hat_noisy = decode(z_tilde)
loss_noisy  = mse(x, x_hat_noisy)          # <-- against the CLEAN x, not x_tilde
print("Denoised: z =", z_tilde, " x_hat =", x_hat_noisy, " MSE =", round(loss_noisy, 5))
# expected: Clean MSE ~ 0.14442, Denoised MSE ~ 0.19207

# ---- sparse autoencoder: KL-divergence sparsity penalty ----
def kl_sparsity(rho, rho_hat):
    return rho * np.log(rho / rho_hat) + (1 - rho) * np.log((1 - rho) / (1 - rho_hat))

rho       = 0.05                                    # target average activation
rho_hats  = np.array([0.03, 0.06, 0.30, 0.08])       # 4 latent units' actual average activation
penalties = kl_sparsity(rho, rho_hats)
for j, (rh, p) in enumerate(zip(rho_hats, penalties), start=1):
    print(f"  unit z{j}: rho_hat={rh:.2f}  KL penalty={p:.5f}")
print("Total sparsity penalty (sum over units):", round(penalties.sum(), 5))