with Ada.Containers; use Ada.Containers;

package body Conv2d_Fp64_Pkg
  with SPARK_Mode => On
is

   function Convolution_Sum_FP64
     (Input           : Feature_Map_Type;
      Filter          : Filter_Bank_Type;
      OC              : Integer;
      OH              : Integer;
      OW              : Integer;
      Input_Channels  : Integer;
      Kernel_Height   : Integer;
      Kernel_Width    : Integer;
      Input_Width     : Integer) return Interfaces.IEEE_Float_64
   is
      use Interfaces;
      Sum         : IEEE_Float_64 := 0.0;
      Input_Idx   : Natural;
      Filter_Idx  : Natural;
      Channel_Stride : constant Integer := Input.Height * Input_Width;
      Filter_OC_Stride : constant Integer :=
        Input_Channels * Kernel_Height * Kernel_Width;
      Filter_IC_Stride : constant Integer :=
        Kernel_Height * Kernel_Width;
   begin
      for IC in 0 .. Input_Channels - 1 loop
         for KH in 0 .. Kernel_Height - 1 loop
            for KW in 0 .. Kernel_Width - 1 loop
               Input_Idx := IC * Channel_Stride
                            + (OH + KH) * Input_Width
                            + (OW + KW);
               Filter_Idx := OC * Filter_OC_Stride
                             + IC * Filter_IC_Stride
                             + KH * Kernel_Width
                             + KW;
               Sum := Sum
                    + IEEE_Float_64_Vectors.Element (Input.Data, Input_Idx)
                    * IEEE_Float_64_Vectors.Element (Filter.Data, Filter_Idx);
            end loop;
         end loop;
      end loop;
      return Sum;
   end Convolution_Sum_FP64;

   procedure Convolve_2d_Fp64
     (Input           : Feature_Map_Type;
      Filter          : Filter_Bank_Type;
      Output          : in out Feature_Map_Type;
      Input_Height    : Integer;
      Input_Width     : Integer;
      Input_Channels  : Integer;
      Kernel_Height   : Integer;
      Kernel_Width    : Integer;
      Output_Channels : Integer)
   is
      use Interfaces;
      Output_Height : constant Integer := Input_Height - Kernel_Height + 1;
      Output_Width  : constant Integer := Input_Width - Kernel_Width + 1;
      Output_Plane  : constant Integer := Output_Height * Output_Width;
      Input_Plane   : constant Integer := Input_Height * Input_Width;
      Filter_OC_Stride : constant Integer :=
        Input_Channels * Kernel_Height * Kernel_Width;
      Filter_IC_Stride : constant Integer :=
        Kernel_Height * Kernel_Width;
      Expected_Length : constant Count_Type :=
        Count_Type (Output_Channels * Output_Plane);
      Sum         : IEEE_Float_64;
      Out_Idx     : Natural;
      Input_Idx   : Natural;
      Filter_Idx  : Natural;
   begin
      for OC in 0 .. Output_Channels - 1 loop
         pragma Loop_Invariant
           (IEEE_Float_64_Vectors.Length (Output.Data) = Expected_Length);
         pragma Loop_Invariant (Output.Channels = Output_Channels);
         pragma Loop_Invariant (Output.Height = Output_Height);
         pragma Loop_Invariant (Output.Width = Output_Width);

         for OH in 0 .. Output_Height - 1 loop
            pragma Loop_Invariant
              (IEEE_Float_64_Vectors.Length (Output.Data) = Expected_Length);
            pragma Loop_Invariant (Output.Channels = Output_Channels);
            pragma Loop_Invariant (Output.Height = Output_Height);
            pragma Loop_Invariant (Output.Width = Output_Width);

            for OW in 0 .. Output_Width - 1 loop
               pragma Loop_Invariant
                 (IEEE_Float_64_Vectors.Length (Output.Data) = Expected_Length);
               pragma Loop_Invariant (Output.Channels = Output_Channels);
               pragma Loop_Invariant (Output.Height = Output_Height);
               pragma Loop_Invariant (Output.Width = Output_Width);

               Sum := 0.0;
               for IC in 0 .. Input_Channels - 1 loop
                  for KH in 0 .. Kernel_Height - 1 loop
                     for KW in 0 .. Kernel_Width - 1 loop
                        Input_Idx := IC * Input_Plane
                                     + (OH + KH) * Input_Width
                                     + (OW + KW);
                        Filter_Idx := OC * Filter_OC_Stride
                                      + IC * Filter_IC_Stride
                                      + KH * Kernel_Width
                                      + KW;
                        Sum := Sum
                             + IEEE_Float_64_Vectors.Element
                                 (Input.Data, Input_Idx)
                             * IEEE_Float_64_Vectors.Element
                                 (Filter.Data, Filter_Idx);
                     end loop;
                  end loop;
               end loop;

               Out_Idx := OC * Output_Plane + OH * Output_Width + OW;
               IEEE_Float_64_Vectors.Replace_Element
                 (Output.Data, Out_Idx, Sum);
            end loop;
         end loop;
      end loop;
   end Convolve_2d_Fp64;

end Conv2d_Fp64_Pkg;