submodule(nf_conv1d_layer) nf_conv1d_layer_submodule

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

  implicit none

contains

  module function conv1d_layer_cons(filters, kernel_size, activation, stride) result(res)
    integer, intent(in) :: filters
    integer, intent(in) :: kernel_size
    class(activation_function), intent(in) :: activation
    integer, intent(in) :: stride
    type(conv1d_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 conv1d_layer_cons

  module subroutine init(self, input_shape)
    class(conv1d_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

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

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

    ! Kernel of shape: filters x channels x kernel_size
    allocate(self % kernel(self % filters, self % channels, self % kernel_size))
    call random_normal(self % kernel)
    self % kernel = self % kernel / self % kernel_size

    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)))
    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(conv1d_layer), intent(in out) :: self
    real, intent(in) :: input(:,:)
    integer :: input_width
    integer :: j, n
    integer :: iws, iwe

    input_width = size(input, dim=2)

    ! Loop over output positions.
    do j = 1, self % width
      ! Compute the input window corresponding to output index j.
      ! In forward: center index = j + half_window, so window = indices j to j+kernel_size-1.
      iws = self % stride * (j-1) + 1
      iwe = min(iws + self % kernel_size - 1, input_width)

      ! For each filter, compute the convolution (inner product over channels and kernel width).
      do concurrent (n = 1:self % filters)
        self % z(n, j) = sum(self % kernel(n,:,1:iwe-iws+1) * input(:,iws:iwe))
      end do

      ! Add the bias for each filter.
      self % z(:,j) = self % z(:,j) + self % biases
    end do

    ! Apply the activation function.
    self % output = self % activation % eval(self % z)
  end subroutine forward

  pure module subroutine backward(self, input, gradient)
    class(conv1d_layer), intent(in out) :: self
    ! 'input' has shape: (channels, input_width)
    ! 'gradient' (dL/dy) has shape: (filters, output_width)
    real, intent(in) :: input(:,:)
    real, intent(in) :: gradient(:,:)

    integer :: input_channels, input_width
    integer :: j, n, k
    integer :: iws, iwe

    ! Local arrays to accumulate gradients.
    real :: gdz(self % filters, self % width)  ! local gradient (dL/dz)
    real :: db_local(self % filters)
    real :: dw_local(self % filters, self % channels, self % kernel_size)

    input_width = size(input, dim=2)

    !--- Compute the local gradient gdz = (dL/dy) * sigma'(z) for each output.
    gdz = gradient * self % activation % eval_prime(self % z)

    !--- Compute bias gradients: db(n) = sum_j gdz(n, j)
    db_local = sum(gdz, dim=2)

    !--- Initialize weight gradient and input gradient accumulators.
    dw_local = 0.0
    self % gradient = 0.0
    
    !--- Accumulate gradients over each output position.
    ! In the forward pass the window for output index j was:
    !   iws = j,  iwe = j + kernel_size - 1.
    do n = 1, self % filters
      do j = 1, self % width
        iws = self % stride * (j-1) + 1
        iwe = min(iws + self % kernel_size - 1, input_width)
        do k = 1, self % channels
          ! Weight gradient: accumulate contribution from the input window.
          dw_local(n,k,1:iwe-iws+1) = dw_local(n,k,1:iwe-iws+1) + input(k,iws:iwe) * gdz(n,j)
          ! Input gradient: propagate gradient back to the input window.
          self % gradient(k,iws:iwe) = self % gradient(k,iws:iwe) + self % kernel(n,k,1:iwe-iws+1) * gdz(n,j)
        end do
      end do
    end do

    !--- Update stored gradients.
    self % dw = self % dw + dw_local
    self % db = self % db + db_local

  end subroutine backward

  pure module function get_num_params(self) result(num_params)
    class(conv1d_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(conv1d_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(conv1d_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_conv1d_layer_submodule