with Ada.Numerics.Generic_Elementary_Functions;

package body Layernorm_Pkg
  with SPARK_Mode => On
is

   package Float_Functions is new Ada.Numerics.Generic_Elementary_Functions
     (Interfaces.IEEE_Float_64);

   ------------------------------------------------------------------
   --  LN_Mean
   ------------------------------------------------------------------
   function LN_Mean
     (Input_Row : Activation_Vector_Type;
      Width     : Integer) return Interfaces.IEEE_Float_64
   is
      use type Interfaces.IEEE_Float_64;
      Sum : Interfaces.IEEE_Float_64 := 0.0;
   begin
      for I in 0 .. Width - 1 loop
         Sum := Sum + IEEE_Float_64_Vectors.Element (Input_Row.Values, I);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width);
      end loop;

      if Width = 0 then
         return 0.0;
      else
         return Sum / Interfaces.IEEE_Float_64 (Width);
      end if;
   end LN_Mean;

   ------------------------------------------------------------------
   --  LN_Variance
   ------------------------------------------------------------------
   function LN_Variance
     (Input_Row : Activation_Vector_Type;
      Width     : Integer;
      Mean      : Interfaces.IEEE_Float_64) return Interfaces.IEEE_Float_64
   is
      use type Interfaces.IEEE_Float_64;
      Sum  : Interfaces.IEEE_Float_64 := 0.0;
      Diff : Interfaces.IEEE_Float_64;
      Result : Interfaces.IEEE_Float_64;
   begin
      for I in 0 .. Width - 1 loop
         Diff := IEEE_Float_64_Vectors.Element (Input_Row.Values, I) - Mean;
         Sum  := Sum + Diff * Diff;
         pragma Loop_Invariant (Sum >= 0.0);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width);
      end loop;

      if Width = 0 then
         return 0.0;
      end if;

      Result := Sum / Interfaces.IEEE_Float_64 (Width);

      if Result < 0.0 then
         return 0.0;
      else
         return Result;
      end if;
   end LN_Variance;

   ------------------------------------------------------------------
   --  LN_Normalised_Element
   ------------------------------------------------------------------
   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
   is
      use type Interfaces.IEEE_Float_64;
      Mean     : constant Interfaces.IEEE_Float_64 :=
        LN_Mean (Input_Row, Width);
      Variance : constant Interfaces.IEEE_Float_64 :=
        LN_Variance (Input_Row, Width, Mean);
      Denom    : Interfaces.IEEE_Float_64;
      X        : constant Interfaces.IEEE_Float_64 :=
        IEEE_Float_64_Vectors.Element (Input_Row.Values, I);
      S        : constant Interfaces.IEEE_Float_64 :=
        IEEE_Float_64_Vectors.Element (Scale.Values, I);
      B        : constant Interfaces.IEEE_Float_64 :=
        IEEE_Float_64_Vectors.Element (Shift.Values, I);
   begin
      Denom := Float_Functions.Sqrt (Variance + Epsilon);
      if Denom = 0.0 then
         return B;
      end if;
      return S * (X - Mean) / Denom + B;
   end LN_Normalised_Element;

   ------------------------------------------------------------------
   --  Layer_Normalise
   ------------------------------------------------------------------
   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)
   is
      use type Interfaces.IEEE_Float_64;

      Sum      : Interfaces.IEEE_Float_64 := 0.0;
      Mean     : Interfaces.IEEE_Float_64;
      Var_Sum  : Interfaces.IEEE_Float_64 := 0.0;
      Variance : Interfaces.IEEE_Float_64;
      Denom    : Interfaces.IEEE_Float_64;
      Diff     : Interfaces.IEEE_Float_64;
      X, S, B  : Interfaces.IEEE_Float_64;
      Val      : Interfaces.IEEE_Float_64;
      Empty    : IEEE_Float_64_Vectors.Vector;
   begin
      --  Pass 1: compute the row mean.
      for I in 0 .. Width - 1 loop
         Sum := Sum + IEEE_Float_64_Vectors.Element (Input_Row.Values, I);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width);
      end loop;
      Mean := Sum / Interfaces.IEEE_Float_64 (Width);

      --  Pass 2: compute the row variance.
      for I in 0 .. Width - 1 loop
         Diff    := IEEE_Float_64_Vectors.Element (Input_Row.Values, I) - Mean;
         Var_Sum := Var_Sum + Diff * Diff;
         pragma Loop_Invariant (Var_Sum >= 0.0);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width);
      end loop;
      Variance := Var_Sum / Interfaces.IEEE_Float_64 (Width);
      if Variance < 0.0 then
         Variance := 0.0;
      end if;

      Denom := Float_Functions.Sqrt (Variance + Epsilon);

      --  Build the output vector field-by-field (no box aggregate).
      Output_Row := (Width  => Width,
                     Values => Empty);

      --  Pass 3: write each normalised element.
      for I in 0 .. Width - 1 loop
         X := IEEE_Float_64_Vectors.Element (Input_Row.Values, I);
         S := IEEE_Float_64_Vectors.Element (Scale.Values,     I);
         B := IEEE_Float_64_Vectors.Element (Shift.Values,     I);

         if Denom = 0.0 then
            Val := B;
         else
            Val := S * (X - Mean) / Denom + B;
         end if;

         IEEE_Float_64_Vectors.Append (Output_Row.Values, Val);

         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Output_Row.Values)) = I + 1);
         pragma Loop_Invariant (Output_Row.Width = Width);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Input_Row.Values)) = Width);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Scale.Values))     = Width);
         pragma Loop_Invariant
           (Natural (IEEE_Float_64_Vectors.Length (Shift.Values))     = Width);
      end loop;
   end Layer_Normalise;

end Layernorm_Pkg;