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

use type Ada.Containers.Count_Type;

package Max_Pool_Pkg
  with SPARK_Mode => On
is

   type Pool_Config_Type is record
      Window_Height                  : Integer;
      Window_Width                   : Integer;
      Stride                         : Integer;
   end record;

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

   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
      Channel_Count                  : Integer;
      Height                         : Integer;
      Width                          : Integer;
      Values                         : 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);

   function Max_Pool
     (Input         : Feature_Map_Type;
      Channel_Count : Integer;
      Input_Height  : Integer;
      Input_Width   : Integer;
      Window_Height : Integer;
      Window_Width  : Integer;
      Stride        : Integer) return Feature_Map_Type
   with
     Pre  => Channel_Count > 0
             and then Input_Height > 0
             and then Input_Width > 0
             and then Window_Height > 0
             and then Window_Width > 0
             and then Window_Height <= Input_Height
             and then Window_Width <= Input_Width
             and then Stride > 0
             and then Input.Channel_Count = Channel_Count
             and then Input.Height = Input_Height
             and then Input.Width = Input_Width
             and then IEEE_Float_64_Vectors.Length (Input.Values) =
                      Ada.Containers.Count_Type
                        (Channel_Count * Input_Height * Input_Width),
     Post => Max_Pool'Result.Channel_Count = Channel_Count
             and then Max_Pool'Result.Height =
                      (Input_Height - Window_Height) / Stride + 1
             and then Max_Pool'Result.Width =
                      (Input_Width - Window_Width) / Stride + 1
             and then IEEE_Float_64_Vectors.Length
                        (Max_Pool'Result.Values) =
                      Ada.Containers.Count_Type
                        (Max_Pool'Result.Channel_Count
                         * Max_Pool'Result.Height
                         * Max_Pool'Result.Width)
             and then (for all IC in 1 .. Channel_Count =>
                        (for all OH in 1 .. Max_Pool'Result.Height =>
                          (for all OW in 1 .. Max_Pool'Result.Width =>
                            (for all KH in 1 .. Window_Height =>
                              (for all KW in 1 .. Window_Width =>
                                Interfaces.">="
                                  (IEEE_Float_64_Vectors.Element
                                     (Max_Pool'Result.Values,
                                      ((IC - 1) * Max_Pool'Result.Height
                                       + (OH - 1)) * Max_Pool'Result.Width
                                      + (OW - 1)),
                                   IEEE_Float_64_Vectors.Element
                                     (Input.Values,
                                      ((IC - 1) * Input_Height
                                       + ((OH - 1) * Stride + KH - 1))
                                      * Input_Width
                                      + ((OW - 1) * Stride + KW - 1))))))));

end Max_Pool_Pkg;