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

use type Ada.Containers.Count_Type;
use type Interfaces.Integer_32;
use type Interfaces.Integer_8;

package body Int8_Matmul_Pkg
  with SPARK_Mode => On
is

   procedure Int8_Matmul
     (M   : Integer;
      N   : Integer;
      K   : Integer;
      A   : Int8_Matrix_Type;
      LDA : Integer;
      B   : Int8_Matrix_Type;
      LDB : Integer;
      LDC : Integer;
      C   : out Int32_Matrix_Type)
   is
      Total_Elements : constant Integer := LDC * N;
      Acc            : Interfaces.Integer_32;
      A_Val          : Interfaces.Integer_32;
      B_Val          : Interfaces.Integer_32;
   begin
      C := (Data              => Integer_32_Vectors.Empty_Vector,
            Leading_Dimension => LDC,
            Row_Count         => M,
            Column_Count      => N);

      for Idx in 1 .. Total_Elements loop
         Integer_32_Vectors.Append (C.Data, 0);
         pragma Loop_Invariant
           (Ada.Containers.Count_Type
              (Integer_32_Vectors.Length (C.Data))
            = Ada.Containers.Count_Type (Idx));
         pragma Loop_Invariant (C.Row_Count = M);
         pragma Loop_Invariant (C.Column_Count = N);
         pragma Loop_Invariant (C.Leading_Dimension = LDC);
      end loop;

      for J in 1 .. N loop
         for I in 1 .. M loop
            Acc := 0;
            for L in 1 .. K loop
               A_Val := Interfaces.Integer_32
                          (Integer_8_Vectors.Element
                             (A.Data, (I - 1) + (L - 1) * LDA));
               B_Val := Interfaces.Integer_32
                          (Integer_8_Vectors.Element
                             (B.Data, (L - 1) + (J - 1) * LDB));
               Acc := Acc + A_Val * B_Val;
               pragma Loop_Invariant
                 (Acc = Sum_Of_Products (A, B, I, J, L, LDA, LDB));
            end loop;

            Integer_32_Vectors.Replace_Element
              (C.Data, (I - 1) + (J - 1) * LDC, Acc);

            pragma Loop_Invariant
              (Ada.Containers.Count_Type
                 (Integer_32_Vectors.Length (C.Data))
               >= Ada.Containers.Count_Type (LDC * N));
            pragma Loop_Invariant (C.Row_Count = M);
            pragma Loop_Invariant (C.Column_Count = N);
            pragma Loop_Invariant (C.Leading_Dimension = LDC);
         end loop;

         pragma Loop_Invariant
           (Ada.Containers.Count_Type
              (Integer_32_Vectors.Length (C.Data))
            >= Ada.Containers.Count_Type (LDC * N));
         pragma Loop_Invariant (C.Row_Count = M);
         pragma Loop_Invariant (C.Column_Count = N);
         pragma Loop_Invariant (C.Leading_Dimension = LDC);
      end loop;
   end Int8_Matmul;

end Int8_Matmul_Pkg;