with SMG_Bit_Keccak; use SMG_Bit_Keccak;
with Ada.Exceptions; use Ada.Exceptions;
with Ada.Text_IO; use Ada.Text_IO;
with Ada.Strings.Fixed; use Ada.Strings.Fixed;
with Interfaces; use Interfaces;

procedure SMG_Bit_Keccak.Test is
  --types
  type Keccak_Perms is ( None, Theta, Rho, Pi, Chi, Iota );
  type Test_Vector  is array( Keccak_Perms ) of State;
  type Test_Round   is array( Round_Index ) of Test_Vector;
  subtype Hexstring is String( 1 .. Z_Length / 4 ); --word as hex string
  subtype Bitstring is String( 1 .. Z_Length ); -- word as binary string
  type Bithex       is array( 0 .. 3 ) of Bit;

  -- helper methods
  procedure HexCharToBit( H : in Character; B: out Bithex) is
  begin
    case H is
      when '0' => B := (0, 0, 0, 0);
      when '1' => B := (0, 0, 0, 1);
      when '2' => B := (0, 0, 1, 0);
      when '3' => B := (0, 0, 1, 1);
      when '4' => B := (0, 1, 0, 0);
      when '5' => B := (0, 1, 0, 1);
      when '6' => B := (0, 1, 1, 0);
      when '7' => B := (0, 1, 1, 1);
      when '8' => B := (1, 0, 0, 0);
      when '9' => B := (1, 0, 0, 1);
      when 'A' => B := (1, 0, 1, 0);
      when 'B' => B := (1, 0, 1, 1);
      when 'C' => B := (1, 1, 0, 0);
      when 'D' => B := (1, 1, 0, 1);
      when 'E' => B := (1, 1, 1, 0);
      when 'F' => B := (1, 1, 1, 1);
      when others => null;
    end case;
  end HexCharToBit;

  function HexToBitword( H: in Hexstring ) return Bitword is
    BW         : Bitword;
    B1, B2     : Bithex;
    PosH, PosB : Natural;
  begin
    -- read the hexstring octet by octet
    for I in 1 .. Z_Length / 8 loop
      PosH := Integer(H'First) + (I - 1) * 2;
      HexCharToBit( H(PosH), B1 );
      HexCharToBit( H(PosH + 1), B2 );

      PosB := Integer(BW'First) + (I - 1) * 8;
      for J in 0 .. 3 loop
        BW ( ZCoord(PosB + J) ) := B1(J);
        BW ( ZCoord(PosB + 4 + J) ) := B2(J);
      end loop;
    end loop;
    return BW;
  end HexToBitword;

  -- prints one bitword as an array of bits
  procedure print_bitword( B: in Bitword ) is
    bstr: Bitstring;
  begin
    for I in ZCoord loop
      if B( I ) > 0 then
        bstr( Bitstring'First + Integer(I) ) := '1';
      else
        bstr( Bitstring'First + Integer(I) ) := '0';
      end if;
    end loop;
    Put(bstr);
  end print_bitword;

  -- prints a keccak state, bitword by bitword
  procedure print_state( S: in State; Title: in String) is
  begin
    Put_Line("---------" & Title & "---------");
    for Y in XYCoord loop
      for X in XYCoord loop
        Put( "S(" & XYCoord'Image(X) & ", " & XYCoord'Image(Y) & ")= ");
        print_bitword( S( X, Y ) ); 
        new_line(1);
      end loop;
    end loop;
  end print_state;

  function read_state(File: in FILE_TYPE; Oct: Positive :=8) return State is
    S: State;
    Line1: String := "0000000000000000 " &
                     "0000000000000000 " &
                     "0000000000000000 " & 
                     "0000000000000000 " &
                     "0000000000000000";
    StartPos, EndPos: Positive;
    Len: Positive := Oct*2;
    HStr: Hexstring;
  begin
    for Y in XYCoord loop
      Line1 := Get_Line(File);
      StartPos := Line1'First;
      EndPos := StartPos + Len-1;

      for X in XYCoord loop
        HStr := Line1( StartPos .. EndPos );
        S( X, Y ) := HexToBitword(HStr);
        StartPos := EndPos + 2;	--one space to skip
        EndPos := StartPos + Len - 1;
      end loop;
    end loop;
    return S;
  end read_state;

  --reads a full test round from specified file (pre-defined format)
  function read_from_file (filename : in String;
                           T        : out Test_Round)
                           return Boolean is 
    file: FILE_TYPE;
    InputMarker: String := "lanes as 64-bit words:";
    octets: Positive := 8;
    RoundNo: Round_Index;
  begin
    -- try to open the input file
    begin
      open(file, In_File, filename);
    exception
      when others =>
        Put_Line(Standard_Error, 
                 "Can not open the file '" & filename & "'. Does it exist?");
        return False;
    end;

  -- find & read input state first
    RoundNo := -1;
    loop
      declare
        Line: String := Get_Line(file);
      begin
        --check if this is test data of any known kind
        if index(Line, InputMarker, 1) > 0 then
          T(0)(None) := read_state(file, octets);
          print_state(T(0)(None), "Read Input State");
        elsif index(Line, "Round ", 1) > 0 then
          RoundNo := RoundNo +1;
        elsif index(Line, "theta", 1) > 0 then
          T(RoundNo)(Theta) := read_state(file, octets);
          if (RoundNo > 0) then
            T(RoundNo)(None) := T(RoundNo-1)(Iota);  -- previous state as input
          end if;
        elsif index(Line, "rho", 1) > 0 then
          T(RoundNo)(Rho) := read_state(file, octets);
        elsif index(Line, "pi", 1) > 0 then
          T(RoundNo)(Pi) := read_state(file, octets);
        elsif index(Line, "chi", 1) > 0 then
          T(RoundNo)(Chi) := read_state(file, octets);
        elsif index(Line, "iota", 1) > 0 then
          T(RoundNo)(Iota) := read_state(file, octets);
        end if;	
        exit when End_Of_File(file);
      end;
    end loop;
    Close(file);
    return True;
  end read_from_file;


  -- performs one single round of Keccak, step by step
  -- each permutation is tested separately
  -- test fails with exception raised at first output not matching expected
  procedure test_one_round(T: Test_Vector; Round: Round_Index) is
    Input: State;
    Expected: State;
    Output: State;
    Test_One_Round_Fail: Exception;
  begin
    Input := T(None);
    for I in Keccak_Perms range Theta..Iota loop
      Expected := T(I);
      case I is
        when Theta  => Output := SMG_Bit_Keccak.Theta(Input);
        when Rho    => Output := SMG_Bit_Keccak.Rho(Input);
        when Pi     => Output := SMG_Bit_Keccak.Pi(Input);
        when Chi    => Output := SMG_Bit_Keccak.Chi(Input);
        when Iota   => Output := SMG_Bit_Keccak.Iota(RC(Round), Input);
        when others => null;
      end case;

      if (Output /= Expected) then
        print_state(Output, "----------real output-------");
        print_state(Expected, "----------expected output--------");
        raise Test_One_Round_Fail;
      else
        Put_Line("PASSED: " & Keccak_Perms'Image(I));
      end if;
      -- get ready for next permutation
      Input := Expected;
    end loop;
  end test_one_round;

  procedure test_bwrotate_left( Input    : in Bitword;
                                N        : Positive;
                                Expected : in Bitword) is
    Output: Bitword;
  begin
    Output := BWRotate_Left( Input, N );
    if Output /= Expected then
      Put_Line("FAIL: test bitword rotate left");
      Put_Line("Output:");
      print_bitword( Output );
      Put_Line("Expected:");
      print_bitword( Expected );
    else
      Put_Line("PASS: test bitword rotate left");
    end if;
  end test_bwrotate_left;

  procedure test_keccak_function(T: in Test_Round) is
    S: State;
  begin
    Put_Line("---Full Keccak Function test---");
    S := Keccak_Function(T(Round_Index'First)(None));
    if S /= T(Round_Index'Last)(Iota) then
      Put_Line("FAILED: full keccak function test");
    else
      Put_Line("PASSED: full keccak function test");
    end if;
  end test_keccak_function;

  procedure test_sponge is
    Bitrate   : constant Keccak_Rate := 1344;
    Input1    : Bitstream( 1 .. 5 ) := (1, 1, 0, 0, 1);
    Input2    : Bitstream( 1 .. 30) := (1, 1, 0, 0, 
                                        1, 0, 1, 0, 
                                        0, 0, 0, 1,
                                        1, 0, 1, 0, 
                                        1, 1, 0, 1, 
                                        1, 1, 1, 0, 
                                        1, 0, 0, 1,
                                        1, 0); 
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    C         : Natural;
    ExpHex1   : constant String := 
              "CB7FFB7CE7572A06C537858A0090FC2888C3C6BA9A3ADAB4"&
              "FE7C9AB4EFE7A1E619B834C843A5A79E23F3F7E314AA597D"&
              "9DAD376E8413A005984D00CF954F62F59EF30B050C99EA64"&
              "E958335DAE684195D439B6E6DFD0E402518B5E7A227C48CF"&
              "239CEA1C391241D7605733A9F4B8F3FFBE74EE45A40730ED"&
              "1E2FDEFCCA941F518708CBB5B6D5A69C30263267B97D7B29"&
              "AC87043880AE43033B1017EFB75C33248E2962892CE69DA8"&
              "BAF1DF4C0902B16C64A1ADD42FF458C94C4D3B0B32711BBA"&
              "22104989982543D1EF1661AFAF2573687D588C81113ED7FA"&
              "F7DDF912021FC03D0E98ACC0200A9F7A0E9629DBA33BA0A3"&
              "C03CCA5A7D3560A6DB589422AC64882EF14A62AD9807B353"&
              "8DEE1548194DBD456F92B568CE76827F41E0FB3C7F25F3A4"&
              "C707AD825B289730FEBDFD22A3E742C6FB7125DE0E38B130"&
              "F3059450CA6185156A7EEE2AB7C8E4709956DC6D5E9F99D5"&
              "0A19473EA7D737AC934815D68C0710235483DB8551FD8756"&
              "45692B4E5E16BB9B1142AE300F5F69F43F0091D534F372E1"&
              "FFC2E522E71003E4D27EF6ACCD36B2756FB5FF02DBF0C96B"&
              "CAE68E7D6427810582F87051590F6FB65D7B948A9C9D6C93"&
              "AF4562367A0AD79109D6F3087C775FE6D60D66B74F8D29FB"&
              "4BA80D0168693A748812EA0CD3CA23854CC84D4E716F4C1A"&
              "A3B340B1DED2F304DFDBACC1D792C8AC9A1426913E3F67DB"& 
              "790FD5CFB77DAA29"; 
    ExpHex2   : constant String :=
              "35F4FBA9D29E833B1DB17CA2077C11B3348C8AF2A29344AE"&
              "6AAA1F63FC4536CE795C54F0359953B97CEA27491691E93E"&
              "E4829EAB388211E6E8BD3EDA74366D0947DFA3D65D127593"&
              "0AFC42884B7324717DCB003D7B3B5C2E92B84F478CC8DBB5"&
              "174EB4BAC6207BD22E56FCC6E5FB11BC598FDBE6208913CE"&
              "34BC03837FDBFCDFF9407D948531B5FC7FFE7029F30E7EDC"&
              "F9282F0A630FA99839776F5EEA485449F62E421552AF9571"; 
    HexStr1   : String( 1 .. ExpHex1'Length );
    Output1   : Bitstream( 1 .. ExpHex1'Length * 4 );
    HexStr2   : String( 1 .. ExpHex2'Length );
    Output2   : Bitstream( 1 .. ExpHex2'Length * 4 );
    Error     : Natural;
    Pos       : Natural;
    HexPos    : Natural;
  begin

  -- test 1
    Put_Line("---sponge test 1---");
    Sponge(Input1, Bitrate, Output1);
    Put_Line("Input is:");
    for I of Input1 loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Put_Line("Output is:");
    for I of Output1 loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Error := 0;
    for I in 1..Output1'Length/4 loop
      Pos := Output1'First + (I-1)*4;
      C := Natural( Output1( Pos ) ) +
           Natural( Output1( Pos + 1 ) ) * 2 +
           Natural( Output1( Pos + 2 ) ) * 4 +
           Natural( Output1( Pos + 3 ) ) * 8;
      HexPos := I + 2 * ( I mod 2 ) - 1; 
			Hexstr1( HexPos ) := Hex(C);
      if Hexstr1( HexPos ) /= ExpHex1( HexPos ) then
        Error := Error + 1;
      end if;
    end loop;
    Put_Line("Expected: ");
    Put_Line(ExpHex1);
    Put_Line("Obtained: ");
    Put_Line(Hexstr1);
    Put_Line("Errors found: " & Natural'Image(Error));

  -- test 2
    Put_Line("---sponge test 2---");
    Sponge(Input2, Bitrate, Output2);
    Put_Line("Input is:");
    for I of Input2 loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Put_Line("Output is:");
    for I of Output2 loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Error := 0;
    for I in 1..Output2'Length/4 loop
      Pos := Output2'First + (I-1)*4;
      C := Natural( Output2( Pos ) ) +
           Natural( Output2( Pos + 1 ) ) * 2 +
           Natural( Output2( Pos + 2 ) ) * 4 +
           Natural( Output2( Pos + 3 ) ) * 8;
      HexPos := I + 2 * ( I mod 2 ) - 1; 
			Hexstr2( HexPos ) := Hex(C);
      if Hexstr2( HexPos ) /= ExpHex2( HexPos ) then
        Error := Error + 1;
      end if;
    end loop;
    Put_Line("Expected: ");
    Put_Line(ExpHex2);
    Put_Line("Obtained: ");
    Put_Line(Hexstr2);
    Put_Line("Errors found: " & Natural'Image(Error));

  end test_sponge;

  -- end of helper methods 

	--variables
  T     : Test_Round;
  BW, E : Bitword;
begin
  BW:=(0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0,
       0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,1,0);
  E:=(0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0,
       0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0, 1,0,0,0);
  test_bwrotate_left(BW, 2, E);

  Put_Line("-----Testing with zero state as input------");
  if (not read_from_file("testvectorszero.txt", T)) then
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  Put_Line("-----Testing with non-zero state as input------");
  if (not read_from_file("testvectorsnonzero.txt", T)) then 
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  -- test Sponge construction
  test_sponge;

end SMG_Bit_Keccak.Test;