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

use type Ada.Containers.Count_Type;

package Int8_Conv2d_Pkg
  with SPARK_Mode => On
is

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

   type Int8_Tensor_3D_Type is record
      Width                          : Integer;
      Height                         : Integer;
      Channels                       : Integer;
      Data                           : Integer_8_Vectors.Vector;
   end record;

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

   type Int8_Tensor_4D_Type is record
      Kernel_Width                   : Integer;
      Kernel_Height                  : Integer;
      Input_Channels                 : Integer;
      Output_Channels                : Integer;
      Data                           : Integer_8_Vectors.Vector;
   end record;

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

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

   type Int32_Tensor_3D_Type is record
      Width                          : Integer;
      Height                         : Integer;
      Channels                       : Integer;
      Data                           : Integer_32_Vectors.Vector;
   end record;

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

   function Output_Cell
     (Y : Int32_Tensor_3D_Type;
      Ox, Oy, Oc : Positive) return Interfaces.Integer_32
   with
      Ghost,
      Pre => Ox <= Y.Width
             and then Oy <= Y.Height
             and then Oc <= Y.Channels;

   function Conv_Sum_At
     (X : Int8_Tensor_3D_Type;
      W : Int8_Tensor_4D_Type;
      Ox, Oy, Oc : Positive) return Interfaces.Integer_32
   with
      Ghost,
      Pre => X.Width > 0
             and then X.Height > 0
             and then X.Channels > 0
             and then W.Kernel_Width > 0
             and then W.Kernel_Height > 0
             and then W.Input_Channels > 0
             and then W.Output_Channels > 0
             and then W.Input_Channels = X.Channels
             and then W.Kernel_Width <= X.Width
             and then W.Kernel_Height <= X.Height
             and then W.Kernel_Width * W.Kernel_Height * X.Channels <= 133_164
             and then Ox <= X.Width - W.Kernel_Width + 1
             and then Oy <= X.Height - W.Kernel_Height + 1
             and then Oc <= W.Output_Channels;

   procedure Int8_Conv2d
     (X : Int8_Tensor_3D_Type;
      W : Int8_Tensor_4D_Type;
      Y : out Int32_Tensor_3D_Type)
   with
      Pre  => X.Width > 0
              and then X.Height > 0
              and then X.Channels > 0
              and then W.Kernel_Width > 0
              and then W.Kernel_Height > 0
              and then W.Input_Channels > 0
              and then W.Output_Channels > 0
              and then W.Kernel_Width <= X.Width
              and then W.Kernel_Height <= X.Height
              and then W.Input_Channels = X.Channels
              and then W.Kernel_Width * W.Kernel_Height * X.Channels <= 133_164,
      Post => Y.Width = X.Width - W.Kernel_Width + 1
              and then Y.Height = X.Height - W.Kernel_Height + 1
              and then Y.Channels = W.Output_Channels
              and then (for all Oc in 1 .. Y.Channels =>
                          (for all Oy in 1 .. Y.Height =>
                             (for all Ox in 1 .. Y.Width =>
                                Interfaces."="
                                  (Output_Cell (Y, Ox, Oy, Oc),
                                   Conv_Sum_At (X, W, Ox, Oy, Oc)))));

end Int8_Conv2d_Pkg;