with Interfaces;
with SPARK.Containers.Formal.Unbounded_Vectors;
with Ada;
with Ada.Containers;

use type Ada.Containers.Count_Type;

package Conv2d_Fp64_Pkg
  with SPARK_Mode => On
is

   package IEEE_Float_64_Vectors is new SPARK.Containers.Formal.Unbounded_Vectors
      (Index_Type   => Natural,
       Element_Type => Interfaces.IEEE_Float_64,
       "="          => Interfaces."=");

   type Feature_Map_Type is record
      Channels                       : Integer;
      Height                         : Integer;
      Width                          : Integer;
      Data                           : IEEE_Float_64_Vectors.Vector;
   end record;

   package Feature_Map_Type_Vectors is new SPARK.Containers.Formal.Unbounded_Vectors
      (Index_Type   => Natural,
       Element_Type => Feature_Map_Type);

   type Filter_Bank_Type is record
      Output_Channels                : Integer;
      Input_Channels                 : Integer;
      Kernel_Height                  : Integer;
      Kernel_Width                   : Integer;
      Data                           : IEEE_Float_64_Vectors.Vector;
   end record;

   package Filter_Bank_Type_Vectors is new SPARK.Containers.Formal.Unbounded_Vectors
      (Index_Type   => Natural,
       Element_Type => Filter_Bank_Type);

   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
   with
     Ghost,
     Pre => Input_Channels > 0
         and then Kernel_Height > 0
         and then Kernel_Width > 0
         and then Input_Width > 0
         and then Input.Channels = Input_Channels
         and then Filter.Input_Channels = Input_Channels
         and then Filter.Kernel_Height = Kernel_Height
         and then Filter.Kernel_Width = Kernel_Width
         and then OC >= 0
         and then OC < Filter.Output_Channels
         and then OH >= 0
         and then OH <= Input.Height - Kernel_Height
         and then OW >= 0
         and then OW <= Input.Width - Kernel_Width;

   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)
   with
     Pre  => Input_Height > 0
          and then Input_Width > 0
          and then Input_Channels > 0
          and then Kernel_Height > 0
          and then Kernel_Width > 0
          and then Output_Channels > 0
          and then Kernel_Height <= Input_Height
          and then Kernel_Width <= Input_Width
          and then Input.Channels = Input_Channels
          and then Input.Height = Input_Height
          and then Input.Width = Input_Width
          and then Natural (IEEE_Float_64_Vectors.Length (Input.Data)) =
                    Input_Channels * Input_Height * Input_Width
          and then Filter.Output_Channels = Output_Channels
          and then Filter.Input_Channels = Input_Channels
          and then Filter.Kernel_Height = Kernel_Height
          and then Filter.Kernel_Width = Kernel_Width
          and then Natural (IEEE_Float_64_Vectors.Length (Filter.Data)) =
                    Output_Channels * Input_Channels
                    * Kernel_Height * Kernel_Width
          and then Output.Channels = Output_Channels
          and then Output.Height = Input_Height - Kernel_Height + 1
          and then Output.Width = Input_Width - Kernel_Width + 1
          and then Natural (IEEE_Float_64_Vectors.Length (Output.Data)) =
                    Output_Channels
                    * (Input_Height - Kernel_Height + 1)
                    * (Input_Width - Kernel_Width + 1),
     Post => Output.Channels = Output_Channels
          and then Output.Height = Input_Height - Kernel_Height + 1
          and then Output.Width = Input_Width - Kernel_Width + 1
          and then Natural (IEEE_Float_64_Vectors.Length (Output.Data)) =
                    Output_Channels
                    * (Input_Height - Kernel_Height + 1)
                    * (Input_Width - Kernel_Width + 1)
          and then (for all OC in 0 .. Output_Channels - 1 =>
                      (for all OH in 0 .. Input_Height - Kernel_Height =>
                         (for all OW in 0 .. Input_Width - Kernel_Width =>
                            Interfaces."="
                              (IEEE_Float_64_Vectors.Element
                                 (Output.Data,
                                  OC
                                  * (Input_Height - Kernel_Height + 1)
                                  * (Input_Width - Kernel_Width + 1)
                                  + OH * (Input_Width - Kernel_Width + 1)
                                  + OW),
                               Convolution_Sum_FP64
                                 (Input, Filter, OC, OH, OW,
                                  Input_Channels, Kernel_Height, Kernel_Width,
                                  Input_Width)))));

end Conv2d_Fp64_Pkg;