package body Int8_Conv2d_Pkg
  with SPARK_Mode => On
is

   function Output_Index
     (Y : Int32_Tensor_3D_Type; Ox, Oy, Oc : Positive) return Natural
   is (((Oc - 1) * Y.Height + (Oy - 1)) * Y.Width + (Ox - 1));

   function Input_Index
     (X : Int8_Tensor_3D_Type; Ix, Iy, Ic : Positive) return Natural
   is (((Ic - 1) * X.Height + (Iy - 1)) * X.Width + (Ix - 1));

   function Weight_Index
     (W : Int8_Tensor_4D_Type; Kx, Ky, Ic, Oc : Positive) return Natural
   is ((((Oc - 1) * W.Input_Channels + (Ic - 1)) * W.Kernel_Height
         + (Ky - 1)) * W.Kernel_Width + (Kx - 1));

   function Output_Cell
     (Y : Int32_Tensor_3D_Type;
      Ox, Oy, Oc : Positive) return Interfaces.Integer_32
   is
   begin
      return Integer_32_Vectors.Element
        (Y.Data, Output_Index (Y, Ox, Oy, Oc));
   end Output_Cell;

   function Conv_Sum_At
     (X : Int8_Tensor_3D_Type;
      W : Int8_Tensor_4D_Type;
      Ox, Oy, Oc : Positive) return Interfaces.Integer_32
   is
      use Interfaces;
      Sum : Interfaces.Integer_32 := 0;
   begin
      for Ic in 1 .. X.Channels loop
         for Ky in 1 .. W.Kernel_Height loop
            for Kx in 1 .. W.Kernel_Width loop
               declare
                  Ix : constant Positive := Ox + Kx - 1;
                  Iy : constant Positive := Oy + Ky - 1;
                  X_Val : constant Interfaces.Integer_32 :=
                    Interfaces.Integer_32
                      (Integer_8_Vectors.Element
                         (X.Data, Input_Index (X, Ix, Iy, Ic)));
                  W_Val : constant Interfaces.Integer_32 :=
                    Interfaces.Integer_32
                      (Integer_8_Vectors.Element
                         (W.Data, Weight_Index (W, Kx, Ky, Ic, Oc)));
               begin
                  Sum := Sum + X_Val * W_Val;
               end;
            end loop;
         end loop;
      end loop;
      return Sum;
   end Conv_Sum_At;

   procedure Int8_Conv2d
     (X : Int8_Tensor_3D_Type;
      W : Int8_Tensor_4D_Type;
      Y : out Int32_Tensor_3D_Type)
   is
      use Interfaces;
      Out_W : constant Positive := X.Width  - W.Kernel_Width  + 1;
      Out_H : constant Positive := X.Height - W.Kernel_Height + 1;
      Out_C : constant Positive := W.Output_Channels;
   begin
      Y := (Width    => Out_W,
            Height   => Out_H,
            Channels => Out_C,
            Data     => Integer_32_Vectors.Empty_Vector);

      for Oc_Idx in 1 .. Out_C loop
         for Oy_Idx in 1 .. Out_H loop
            for Ox_Idx in 1 .. Out_W loop
               declare
                  Sum : Interfaces.Integer_32 := 0;
               begin
                  for Ic in 1 .. X.Channels loop
                     for Ky in 1 .. W.Kernel_Height loop
                        for Kx in 1 .. W.Kernel_Width loop
                           declare
                              Ix : constant Positive := Ox_Idx + Kx - 1;
                              Iy : constant Positive := Oy_Idx + Ky - 1;
                              X_Val : constant Interfaces.Integer_32 :=
                                Interfaces.Integer_32
                                  (Integer_8_Vectors.Element
                                     (X.Data,
                                      Input_Index (X, Ix, Iy, Ic)));
                              W_Val : constant Interfaces.Integer_32 :=
                                Interfaces.Integer_32
                                  (Integer_8_Vectors.Element
                                     (W.Data,
                                      Weight_Index
                                        (W, Kx, Ky, Ic, Oc_Idx)));
                           begin
                              Sum := Sum + X_Val * W_Val;
                           end;
                        end loop;
                     end loop;
                  end loop;
                  Integer_32_Vectors.Append (Y.Data, Sum);
               end;
            end loop;
         end loop;
      end loop;
   end Int8_Conv2d;

end Int8_Conv2d_Pkg;