"""
kiri/backend/cpu_conv.py
Naive but correct im2col-based Conv2D for the CPU backend.
"""

import numpy as np


def conv2d_forward_backward(x_t, w_t, b_t, stride, padding):
    from kiri.tensor import _CPUTensor

    x = x_t.data  # (N, C_in, H, W)
    w = w_t.data  # (C_out, C_in, kH, kW)
    b = b_t.data if b_t is not None else None

    N, C_in, H, W = x.shape
    C_out, _, kH, kW = w.shape
    H_out = (H + 2*padding - kH) // stride + 1
    W_out = (W + 2*padding - kW) // stride + 1

    x_pad = np.pad(x, ((0,0),(0,0),(padding,padding),(padding,padding)))

    col = _im2col(x_pad, kH, kW, stride, H_out, W_out)  # (N, C_in*kH*kW, H_out*W_out)
    w_flat = w.reshape(C_out, -1)                         # (C_out, C_in*kH*kW)
    out_flat = w_flat @ col                               # (C_out, N*H_out*W_out)? No...

    # correct: for each sample
    out = np.zeros((N, C_out, H_out, W_out), dtype=np.float32)
    for n in range(N):
        col_n = _im2col(x_pad[n:n+1], kH, kW, stride, H_out, W_out)  # (1, C_in*kH*kW, H_out*W_out)
        col_n = col_n[0]  # (C_in*kH*kW, H_out*W_out)
        out[n] = (w_flat @ col_n).reshape(C_out, H_out, W_out)

    if b is not None:
        out += b.reshape(1, C_out, 1, 1)

    parents = (x_t, w_t) + ((b_t,) if b_t is not None else ())
    result = _CPUTensor(out, parents, "conv2d")

    def _backward():
        dout = result.grad  # (N, C_out, H_out, W_out)
        dx_pad = np.zeros_like(x_pad)
        dw = np.zeros_like(w)
        if b is not None:
            b_t.grad += dout.sum(axis=(0,2,3))

        for n in range(N):
            col_n = _im2col(x_pad[n:n+1], kH, kW, stride, H_out, W_out)[0]
            dout_n = dout[n].reshape(C_out, -1)  # (C_out, H_out*W_out)
            dw += (dout_n @ col_n.T).reshape(w.shape)
            dcol_n = w_flat.T @ dout_n  # (C_in*kH*kW, H_out*W_out)
            dx_pad[n:n+1] += _col2im(dcol_n[np.newaxis], x_pad.shape[2:],
                                      kH, kW, stride, H_out, W_out)

        if padding > 0:
            x_t.grad += dx_pad[:, :, padding:-padding, padding:-padding]
        else:
            x_t.grad += dx_pad
        w_t.grad += dw

    result._backward = _backward
    return result


def _im2col(x_pad, kH, kW, stride, H_out, W_out):
    N, C, _, _ = x_pad.shape
    col = np.zeros((N, C * kH * kW, H_out * W_out), dtype=np.float32)
    for i in range(kH):
        for j in range(kW):
            row_start = i * kW + j
            col[:, row_start*C:(row_start+1)*C, :] = \
                x_pad[:, :,
                      i:i+H_out*stride:stride,
                      j:j+W_out*stride:stride].reshape(N, C, -1)
    return col


def _col2im(col, input_shape, kH, kW, stride, H_out, W_out):
    H, W = input_shape
    N = col.shape[0]
    x_pad = np.zeros((N, col.shape[1]//(kH*kW), H, W), dtype=np.float32)
    C = x_pad.shape[1]
    for i in range(kH):
        for j in range(kW):
            row_start = i * kW + j
            x_pad[:, :,
                  i:i+H_out*stride:stride,
                  j:j+W_out*stride:stride] += \
                col[:, row_start*C:(row_start+1)*C, :].reshape(N, C, H_out, W_out)
    return x_pad
