import numpy as np

# ==============================================================
# Part 1 -- Contractive penalty: Frobenius norm of the encoder Jacobian
#           (reproduces the lecture's linear-encoder example: 0.29)
# ==============================================================
We = np.array([0.4, 0.3, -0.2])          # same linear encoder as Lecture 23

def contractive_penalty_linear(W):
    """For a LINEAR encoder z = Wx + b, the Jacobian is just W itself
    (constant, independent of x), so ||J||_F^2 = sum(W_ji^2)."""
    return np.sum(W ** 2)

def contractive_penalty_sigmoid(W, z):
    """For a sigmoid encoder z = sigmoid(Wx + b):
    ||J||_F^2 = sum_j (z_j(1-z_j))^2 * sum_i W_ji^2   (per latent unit j)."""
    local_grad_sq = (z * (1 - z)) ** 2          # elementwise sigmoid'(z)^2
    row_norms_sq  = np.sum(W ** 2, axis=-1)     # sum_i W_ji^2 per row j
    return np.sum(local_grad_sq * row_norms_sq)

jf_linear = contractive_penalty_linear(We)
print("Linear-encoder contractive penalty ||J_f||_F^2 =", round(jf_linear, 5))
# expected: 0.29

# same weights, imagining instead a sigmoid encoder currently outputting z=0.75
z_example  = np.array([0.75])
jf_sigmoid = contractive_penalty_sigmoid(We.reshape(1, -1), z_example)
print("Sigmoid-encoder contractive penalty at z=0.75:", round(jf_sigmoid, 5))

# ==============================================================
# Part 2 -- Convolutional autoencoder skeleton (Keras), MNIST-shaped input
#           Full training belongs on Colab/GPU -- see the lecture's box.
# ==============================================================
def build_conv_autoencoder():
    from tensorflow.keras import layers, models
    inp = layers.Input(shape=(28, 28, 1))

    # ---- encoder: Conv2D with stride to shrink spatial size ----
    x = layers.Conv2D(16, 3, activation='relu', padding='same', strides=2)(inp)   # 28->14
    x = layers.Conv2D(8,  3, activation='relu', padding='same', strides=2)(x)     # 14->7
    encoded = x                                                                    # bottleneck feature map

    # ---- decoder: Conv2DTranspose ("deconvolution") to upsample back ----
    x = layers.Conv2DTranspose(8,  3, activation='relu', padding='same', strides=2)(encoded)  # 7->14
    x = layers.Conv2DTranspose(16, 3, activation='relu', padding='same', strides=2)(x)         # 14->28
    decoded = layers.Conv2D(1, 3, activation='sigmoid', padding='same')(x)                     # 28x28x1

    autoencoder = models.Model(inp, decoded)
    autoencoder.compile(optimizer='adam', loss='binary_crossentropy')
    return autoencoder

if __name__ == "__main__":
    try:
        model = build_conv_autoencoder()
        model.summary()
    except ImportError:
        print("TensorFlow/Keras not installed -- run this part on Colab (see the lecture's resource link).")