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

use type Ada.Containers.Count_Type;

package Layernorm_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 Activation_Vector_Type is record
      Width                          : Integer;
      Values                         : IEEE_Float_64_Vectors.Vector;
   end record;

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

   type Parameter_Vector_Type is record
      Width                          : Integer;
      Values                         : IEEE_Float_64_Vectors.Vector;
   end record;

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

   function LN_Mean
     (Input_Row : Activation_Vector_Type;
      Width     : Integer) return Interfaces.IEEE_Float_64
   with
     Ghost,
     Pre => Width > 0
            and then Input_Row.Width = Width
            and then Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width;

   function LN_Variance
     (Input_Row : Activation_Vector_Type;
      Width     : Integer;
      Mean      : Interfaces.IEEE_Float_64) return Interfaces.IEEE_Float_64
   with
     Ghost,
     Pre  => Width > 0
             and then Input_Row.Width = Width
             and then Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width,
     Post => Interfaces.">=" (LN_Variance'Result, 0.0);

   function LN_Normalised_Element
     (Input_Row : Activation_Vector_Type;
      Scale     : Parameter_Vector_Type;
      Shift     : Parameter_Vector_Type;
      Width     : Integer;
      Epsilon   : Interfaces.IEEE_Float_64;
      I         : Natural) return Interfaces.IEEE_Float_64
   with
     Ghost,
     Pre => Width > 0
            and then Interfaces.">" (Epsilon, 0.0)
            and then Input_Row.Width = Width
            and then Scale.Width = Width
            and then Shift.Width = Width
            and then Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width
            and then Natural (IEEE_Float_64_Vectors.Length (Scale.Values)) = Width
            and then Natural (IEEE_Float_64_Vectors.Length (Shift.Values)) = Width
            and then I < Width;

   procedure Layer_Normalise
     (Width      : Integer;
      Input_Row  : Activation_Vector_Type;
      Scale      : Parameter_Vector_Type;
      Shift      : Parameter_Vector_Type;
      Epsilon    : Interfaces.IEEE_Float_64;
      Output_Row : out Activation_Vector_Type)
   with
     Pre  => Width > 0
             and then Interfaces.">" (Epsilon, 0.0)
             and then Input_Row.Width = Width
             and then Scale.Width = Width
             and then Shift.Width = Width
             and then Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width
             and then Natural (IEEE_Float_64_Vectors.Length (Scale.Values)) = Width
             and then Natural (IEEE_Float_64_Vectors.Length (Shift.Values)) = Width,
     Post => Output_Row.Width = Width
             and then Natural (IEEE_Float_64_Vectors.Length (Output_Row.Values)) = Width
             and then (for all I in 0 .. Width - 1 =>
                         Interfaces."=" (
                           IEEE_Float_64_Vectors.Element (Output_Row.Values, I),
                           LN_Normalised_Element
                             (Input_Row, Scale, Shift, Width, Epsilon, I)));

end Layernorm_Pkg;