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

use type Ada.Containers.Count_Type;

package body Max_Pool_Pkg
  with SPARK_Mode => On
is

   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
   is
      Output_Height : constant Integer :=
        (Input_Height - Window_Height) / Stride + 1;
      Output_Width  : constant Integer :=
        (Input_Width - Window_Width) / Stride + 1;
      Total_Size    : constant Integer :=
        Channel_Count * Output_Height * Output_Width;
      Result        : Feature_Map_Type :=
        (Channel_Count => Channel_Count,
         Height        => Output_Height,
         Width         => Output_Width,
         Values        => IEEE_Float_64_Vectors.Empty_Vector);
   begin
      for I in 1 .. Total_Size loop
         IEEE_Float_64_Vectors.Append (Result.Values, 0.0);
         pragma Loop_Invariant
           (IEEE_Float_64_Vectors.Length (Result.Values) =
              Ada.Containers.Count_Type (I));
      end loop;

      for IC in 1 .. Channel_Count loop
         for OH in 1 .. Output_Height loop
            for OW in 1 .. Output_Width loop
               declare
                  Out_Idx     : constant Natural :=
                    ((IC - 1) * Output_Height + (OH - 1)) * Output_Width
                    + (OW - 1);
                  Base_In_Idx : constant Natural :=
                    ((IC - 1) * Input_Height + (OH - 1) * Stride)
                    * Input_Width
                    + (OW - 1) * Stride;
                  Max_Val     : Interfaces.IEEE_Float_64 :=
                    IEEE_Float_64_Vectors.Element
                      (Input.Values, Base_In_Idx);
               begin
                  for KH in 1 .. Window_Height loop
                     for KW in 1 .. Window_Width loop
                        declare
                           In_Idx : constant Natural :=
                             ((IC - 1) * Input_Height
                              + ((OH - 1) * Stride + KH - 1))
                             * Input_Width
                             + ((OW - 1) * Stride + KW - 1);
                           Val    : constant Interfaces.IEEE_Float_64 :=
                             IEEE_Float_64_Vectors.Element
                               (Input.Values, In_Idx);
                        begin
                           if Interfaces.">" (Val, Max_Val) then
                              Max_Val := Val;
                           end if;
                        end;
                     end loop;
                  end loop;
                  IEEE_Float_64_Vectors.Replace_Element
                    (Result.Values, Out_Idx, Max_Val);
               end;
            end loop;
         end loop;
      end loop;

      return Result;
   end Max_Pool;

end Max_Pool_Pkg;