submodule(nf_conv2d_layer) nf_conv2d_layer_submodule

  use nf_activation, only: activation_function
  use nf_random, only: random_normal

  implicit none

contains

  module function conv2d_layer_cons(filters, kernel_size, activation, stride) result(res)
    implicit none
    integer, intent(in) :: filters
    integer, intent(in) :: kernel_size
    class(activation_function), intent(in) :: activation
    integer, intent(in) :: stride(:)
    type(conv2d_layer) :: res

    res % kernel_size = kernel_size
    res % filters = filters
    res % activation_name = activation % get_name()
    res % stride = stride
    allocate( res % activation, source = activation )

  end function conv2d_layer_cons


  module subroutine init(self, input_shape)
    implicit none
    class(conv2d_layer), intent(in out) :: self
    integer, intent(in) :: input_shape(:)

    self % channels = input_shape(1)

    self % width = (input_shape(2) - self % kernel_size) / self % stride(1) + 1
    if (mod(input_shape(2) - self % kernel_size , self % stride(1)) /= 0) self % width = self % width + 1

    self % height = (input_shape(3) - self % kernel_size) / self % stride(2) + 1
    if (mod(input_shape(3) - self % kernel_size , self % stride(2)) /= 0) self % height = self % height + 1

    ! Output of shape filters x width x height
    allocate(self % output(self % filters, self % width, self % height))
    self % output = 0

    ! Kernel of shape filters x channels x width x height
    allocate(self % kernel(self % filters, self % channels, &
                           self % kernel_size, self % kernel_size))

    ! Initialize the kernel with random values with a normal distribution.
    call random_normal(self % kernel)
    self % kernel = self % kernel / self % kernel_size**2

    allocate(self % biases(self % filters))
    self % biases = 0

    allocate(self % z, mold=self % output)
    self % z = 0

    allocate(self % gradient(input_shape(1), input_shape(2), input_shape(3)))
    self % gradient = 0

    allocate(self % dw, mold=self % kernel)
    self % dw = 0

    allocate(self % db, mold=self % biases)
    self % db = 0

  end subroutine init


  pure module subroutine forward(self, input)
    class(conv2d_layer), intent(in out) :: self
    real, intent(in) :: input(:,:,:)
    integer :: input_width, input_height, input_channels
    integer :: istart, iend
    integer :: jstart, jend
    integer :: i, j, n
    integer :: iws, iwe, jws, jwe
    integer :: half_window

    ! Input dimensions are channels x width x height
    input_channels = size(input, dim=1)
    input_width = size(input, dim=2)
    input_height = size(input, dim=3)

    ! Half-window is 1 for window size 3; 2 for window size 5; etc.
    half_window = self % kernel_size / 2

    ! Determine the start and end indices for the width and height dimensions
    ! of the input that correspond to the center of each window.
    istart = half_window + 1 ! TODO kernel_width
    jstart = half_window + 1 ! TODO kernel_height

    convolution: do concurrent(i = 1:self % width, j = 1:self % height)

      ! Start and end indices of the input data on the filter window
      ! iws and jws are also coincidentally the indices of the output matrix
      iws = istart + self % stride(1) * (i-1) - half_window ! TODO kernel_width
      iwe = min(iws + 2*half_window, input_width)           ! TODO kernel_width

      jws = jstart + self % stride(2) * (j-1) - half_window ! TODO kernel_height
      jwe = min(jws + 2*half_window, input_height)          ! TODO kernel_height

      ! Compute the inner tensor product, sum(w_ij * x_ij), for each filter.
      do concurrent(n = 1:self % filters)
        self % z(n,i,j) = sum(self % kernel(n,:,1:iwe-iws+1,1:jwe-jws+1) * input(:,iws:iwe,jws:jwe))
      end do

      ! Add bias to the inner product.
      self % z(:,i,j) = self % z(:,i,j) + self % biases

    end do convolution

    ! Activate
    self % output = self % activation % eval(self % z)

  end subroutine forward


  pure module subroutine backward(self, input, gradient)
    class(conv2d_layer), intent(in out) :: self
    real, intent(in) :: input(:,:,:)
    real, intent(in) :: gradient(:,:,:)
    real :: db(self % filters)
    real :: dw(self % filters, self % channels, self % kernel_size, self % kernel_size)
    real :: gdz(self % filters, size(input, 2), size(input, 3))
    integer :: i, j, k, n
    integer :: input_channels, input_width, input_height
    integer :: istart, iend
    integer :: jstart, jend
    integer :: iws, iwe, jws, jwe
    integer :: half_window

    ! Input dimensions are channels x width x height.
    ! Input frame goes from 1 to input_width and from 1 to input_height.
    input_channels = size(input, dim=1)
    input_width = size(input, dim=2)
    input_height = size(input, dim=3)

    ! Half-window is 1 for window size 3; 2 for window size 5; etc.
    half_window = self % kernel_size / 2

    ! Determine the start and end indices for the width and height dimensions
    ! of the input that correspond to the center of each window.
    istart = half_window + 1 ! TODO kernel_width
    jstart = half_window + 1 ! TODO kernel_height
    iend = input_width - istart + 1
    jend = input_height - jstart + 1

    ! z = w .inner. x + b
    ! gdz = dL/dy * sigma'(z)
    gdz = 0
    gdz(:,istart:iend,jstart:jend) = gradient * self % activation % eval_prime(self % z)

    ! dL/db = sum(dL/dy * sigma'(z))
    do concurrent (n = 1:self % filters)
      db(n) = sum(gdz(n,:,:))
    end do

    dw = 0
    self % gradient = 0
    do concurrent( &
      n = 1:self % filters, &
      k = 1:self % channels, &
      i = 1:self % width, &
      j = 1:self % height &
    )
      ! Start and end indices of the input data on the filter window
      iws = istart + self % stride(1) * (i-1) - half_window ! TODO kernel_width
      iwe = min(iws + 2*half_window, input_width)           ! TODO kernel_width

      jws = jstart + self % stride(2) * (j-1) - half_window ! TODO kernel_height
      jwe = min(jws + 2*half_window, input_height)          ! TODO kernel_height

      ! dL/dw = sum(gdz * x)
      dw(n,k,:,:) = dw(n,k,:,:) + input(k,iws:iwe,jws:jwe) * gdz(n,i,j)

      ! dL/dx = sum(gdz * w)
      self % gradient(k,iws:iwe,jws:jwe) = self % gradient(k,iws:iwe,jws:jwe) &
        + gdz(n,i,j) * self % kernel(n,k,:,:)

    end do

    self % dw = self % dw + dw
    self % db = self % db + db

  end subroutine backward


  pure module function get_num_params(self) result(num_params)
    class(conv2d_layer), intent(in) :: self
    integer :: num_params
    num_params = product(shape(self % kernel)) + size(self % biases)
  end function get_num_params

  
  module subroutine get_params_ptr(self, w_ptr, b_ptr)
    class(conv2d_layer), intent(in), target :: self
    real, pointer, intent(out) :: w_ptr(:)
    real, pointer, intent(out) :: b_ptr(:)
    w_ptr(1:size(self % kernel)) => self % kernel
    b_ptr => self % biases
  end subroutine get_params_ptr


  module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
    class(conv2d_layer), intent(in), target :: self
    real, pointer, intent(out) :: dw_ptr(:)
    real, pointer, intent(out) :: db_ptr(:)
    dw_ptr(1:size(self % dw)) => self % dw
    db_ptr => self % db
  end subroutine get_gradients_ptr

end submodule nf_conv2d_layer_submodule