"""
Lecture 13 -- Convolution, Padding, Stride & Pooling
Reproduces: the output-size formula (with/without padding), the
multi-channel conv parameter-count formula, the exact 4x4 max-pooling
worked example, and Sobel Sx/Sy/|grad I| on the worked neighborhood.
"""
import numpy as np

# ---- 1. Output-size formula: W' = (W - K + 2P)/S + 1 ----
def output_size(W, K, S, P=0):
    num = W - K + 2 * P
    if num % S != 0:
        return None  # cannot scan the whole image
    return num // S + 1

for K, S in [(3, 1), (3, 2), (3, 3)]:
    out = output_size(7, K, S)
    print(f"W=7, K={K}, S={S}, P=0 -> {out if out else 'INVALID (does not fit)'}")
print(f"W=7, K=3, S=1, P=1 (same padding) -> {output_size(7, 3, 1, 1)}\n")

# ---- 2. Multi-channel convolution: output shape + parameter count ----
def conv_output(H, W, D_in, K, S, P, D_k):
    Hp, Wp = output_size(H, K, S, P), output_size(W, K, S, P)
    params = (K * K * D_in * D_k) + D_k
    return (Hp, Wp, D_k), params

shape, params = conv_output(28, 28, 3, K=3, S=1, P=1, D_k=15)
print(f"28x28x3 input, 3x3 filter, P=1, 15 filters -> {shape}, {params} params\n")

# ---- 3. Max pooling: exact 4x4 worked example ----
def max_pool(mat, size=2, stride=2):
    mat = np.array(mat)
    out_h = (mat.shape[0] - size)//stride + 1
    out_w = (mat.shape[1] - size)//stride + 1
    out = np.zeros((out_h, out_w))
    for r in range(out_h):
        for c in range(out_w):
            out[r, c] = mat[r*stride:r*stride+size, c*stride:c*stride+size].max()
    return out

X = [[1, 3, 2, 0], [4, 6, 1, 2], [0, 1, 3, 5], [2, 2, 2, 4]]
pooled = max_pool(X)
print("Max pooling 4x4 -> 2x2:\n", pooled)
assert np.array_equal(pooled, [[6, 2], [2, 5]]), "does not match worked example!"

# ---- 4. Sobel edge detection ----
def conv2d_valid(img, kernel):
    img, kernel = np.array(img, float), np.array(kernel, float)
    kh, kw = kernel.shape
    H, W = img.shape
    out = np.zeros((H - kh + 1, W - kw + 1))
    for r in range(out.shape[0]):
        for c in range(out.shape[1]):
            out[r, c] = np.sum(img[r:r+kh, c:c+kw] * kernel)
    return out

Gx = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
Gy = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
N = [[10, 10, 10], [10, 50, 80], [10, 50, 80]]

Sx = conv2d_valid(N, Gx)[0, 0]
Sy = conv2d_valid(N, Gy)[0, 0]
mag = np.sqrt(Sx**2 + Sy**2)
theta = np.degrees(np.arctan2(Sy, Sx))
print(f"\nSobel on N: Sx={Sx:.0f}, Sy={Sy:.0f}, |grad I|={mag:.2f}, theta={theta:.1f} deg")
print("(Course slide rounds |grad I| to ~258.8.)")

# Slide Gx across a 5x5 vertical-edge image (matches the interactive demo)
img5 = [[10, 10, 10, 80, 80]] * 5
print("\nSobel-Gx response over the 5x5 vertical-edge image (3x3 output):\n",
      conv2d_valid(img5, Gx))