# Constantine
# Copyright (c) 2018-2019    Status Research & Development GmbH
# Copyright (c) 2020-Present Mamy André-Ratsimbazafy
# Licensed and distributed under either of
#   * MIT license (license terms in the root directory or at http://opensource.org/licenses/MIT).
#   * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
# at your option. This file may not be copied, modified, or distributed except according to those terms.

import
  # Standard library
  std/[tables, unittest, times],
  # Internal
  ../constantine/[arithmetic, primitives],
  ../constantine/io/[io_fields],
  ../constantine/config/[curves, common],
  # Test utilities
  ../helpers/prng_unsafe


const Iters = 128

var rng: RngState
let seed = uint32(getTime().toUnix() and (1'i64 shl 32 - 1)) # unixTime mod 2^32
rng.seed(seed)
echo "\n------------------------------------------------------\n"
echo "test_finite_fields_sqrt xoshiro512** seed: ", seed

static: doAssert defined(testingCurves), "This modules requires the -d:testingCurves compile option"

proc exhaustiveCheck_p3mod4(C: static Curve, modulus: static int) =
  test "Exhaustive square root check for p ≡ 3 (mod 4) on " & $Curve(C):
    var squares_to_roots: Table[uint16, set[uint16]]

    # Create all squares
    # -------------------------
    for i in 0'u16 ..< modulus:
      var a{.noInit.}: Fp[C]
      a.fromUint(i)

      a.square()

      var r_bytes: array[8, byte]
      r_bytes.exportRawUint(a, cpuEndian)
      let r = uint16(cast[uint64](r_bytes))

      squares_to_roots.mgetOrPut(r, default(set[uint16])).incl(i)

    # From Euler's criterion
    # there is exactly (p-1)/2 squares in 𝔽p* (without 0)
    # and so (p-1)/2 + 1 in 𝔽p (with 0)
    check: squares_to_roots.len == (modulus-1) div 2 + 1

    # Check squares
    # -------------------------
    for i in 0'u16 ..< modulus:
      var a{.noInit.}: Fp[C]
      a.fromUint(i)

      if i in squares_to_roots:
        var a2 = a
        check:
          bool a.isSquare()
          bool a.sqrt_if_square()

        # 2 different code paths have the same result
        # (despite 2 square roots existing per square)
        a2.sqrt()
        check: bool(a == a2)

        var r_bytes: array[8, byte]
        r_bytes.exportRawUint(a, cpuEndian)
        let r = uint16(cast[uint64](r_bytes))

        # r is one of the 2 square roots of `i`
        check: r in squares_to_roots[i]

      else:
        let a2 = a

        check:
          bool not a.isSquare()
          bool not a.sqrt_if_square()
          bool (a == a2) # a shouldn't be modified

proc randomSqrtCheck_p3mod4(C: static Curve) =
  template testImpl(a: untyped): untyped {.dirty.} =
    var na{.noInit.}: Fp[C]
    na.neg(a)

    var a2 = a
    var na2 = na
    a2.square()
    na2.square()
    check:
      bool a2 == na2
      bool a2.isSquare()

    var r, s = a2
    r.sqrt()
    let ok = s.sqrt_if_square()
    check:
      bool ok
      bool(r == s)
      bool(r == a or r == na)

  test "Random square root check for p ≡ 3 (mod 4) on " & $Curve(C):
    for _ in 0 ..< Iters:
      let a = rng.random_unsafe(Fp[C])
      testImpl(a)

    for _ in 0 ..< Iters:
      let a = rng.randomHighHammingWeight(Fp[C])
      testImpl(a)

    for _ in 0 ..< Iters:
      let a = rng.random_long01Seq(Fp[C])
      testImpl(a)

proc main() =
  suite "Modular square root" & " [" & $WordBitwidth & "-bit mode]":
    exhaustiveCheck_p3mod4 Fake103, 103
    exhaustiveCheck_p3mod4 Fake10007, 10007
    exhaustiveCheck_p3mod4 Fake65519, 65519
    randomSqrtCheck_p3mod4 Mersenne61
    randomSqrtCheck_p3mod4 Mersenne127
    randomSqrtCheck_p3mod4 BN254_Nogami
    randomSqrtCheck_p3mod4 BN254_Snarks
    randomSqrtCheck_p3mod4 P256
    randomSqrtCheck_p3mod4 Secp256k1
    randomSqrtCheck_p3mod4 BLS12_381
    randomSqrtCheck_p3mod4 BN446
    randomSqrtCheck_p3mod4 FKM12_447
    randomSqrtCheck_p3mod4 BLS12_461
    randomSqrtCheck_p3mod4 BN462

main()