/*
 *	UCW Library -- Random numbers: Testing
 *
 *	(c) 2020--2022 Pavel Charvat <pchar@ucw.cz>
 *
 *	This software may be freely distributed and used according to the terms
 *	of the GNU Lesser General Public License.
 */

#include <ucw/lib.h>
#include <ucw/random.h>
#include <ucw/opt.h>
#include <ucw/time.h>

#include <errno.h>
#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>

static int o_test_fast;
static int o_test_strong;

static double o_bench_scale = 1.0;
static int o_bench_only_safe;
static int o_bench_all;
static int o_bench_libs;
static int o_bench_legacy;
static int o_bench_fast;
static int o_bench_strong;

static struct opt_section options = {
  OPT_ITEMS {
    OPT_HELP("Usage: random-test <options>"),
    OPT_HELP(""),
    OPT_HELP("Options:"),
    OPT_HELP_OPTION,
    OPT_BOOL(0, "test-fast", o_test_fast, 0, "\tTest fast random generator"),
    OPT_BOOL(0, "test-strong", o_test_strong, 0, "\tTest strong random generator"),
    OPT_BOOL(0, "bench-all", o_bench_all, 0, "\tBench everything"),
    OPT_BOOL(0, "bench-libs", o_bench_libs, 0, "\tBench libc"),
    OPT_BOOL(0, "bench-legacy", o_bench_legacy, 0, "\tBench legacy interface"),
    OPT_BOOL(0, "bench-fast", o_bench_fast, 0, "\tBench fast random generator"),
    OPT_BOOL(0, "bench-strong", o_bench_strong, 0, "\tBench strong random generator"),
    OPT_DOUBLE(0, "bench-scale", o_bench_scale, OPT_REQUIRED_VALUE, "<scale>\tIncrease/decrease the length of benchmark (default: 1.0)"),
    OPT_BOOL(0, "bench-only-safe", o_bench_only_safe, 0, "\tDon't benchmark too verbose or risky functions"),
    OPT_END,
  }
};

#define TEST(x) do { if (unlikely(!(x))) test_failed(#x, __LINE__); } while (0)
static NONRET void test_failed(const char *check, uint line)
{
  msg(L_FATAL, "TEST: Failed at line %u: TEST(%s)", line, check);
  abort();
}

static void test_fast(void)
{
  struct fastrand *r = fastrand_new();
  fastrand_gen_seed(r);
  for (uint i = 0; i < 5000; i++)
    {
      u64 x, y;

      TEST((x = fastrand_max(r, 33)) <= 32);
      TEST((y = fastrand_bits(r, x)) < (1LLU << x));

      TEST((x = fastrand_max(r, 65)) <= 64);
      TEST((y = fastrand_bits_u64(r, x)) <= ((x == 64) ? UINT64_MAX : (1LLU << x) - 1));

      x = fastrand_u32(r) ? : 1;
      TEST((y = fastrand_max(r, x)) < x);
      x = fastrand_u64(r) ? : 1;
      TEST((y = fastrand_max_u64(r, x)) < x);

      double d;
      d = fastrand_double(r);
      TEST(d >= 0.0 - DBL_EPSILON && d < 1.0 + DBL_EPSILON);
    }
  fastrand_delete(r);
  printf("OK\n");
}

static void test_strong(void)
{
  byte buf[100];
  for (uint j = 0; j < 2; j++)
    {
      struct strongrand *r = strongrand_std_open(STRONGRAND_STD_URANDOM, (j == 0) ? 0 : 128);
      for (uint i = 1; i <= 100; i++)
	{
	  strongrand_mem(r, buf, i);
	  if (i % 10 == 0)
	    strongrand_reset(r);
	}
      strongrand_close(r);
    }
  printf("OK\n");
}

#define BENCH(func) BENCH_SLOW(1, func)
#define BENCH_SLOW(mult, func) \
  do { \
    timestamp_t bench_timer; \
    init_timer(&bench_timer); \
    for (uint bench_i = 50000000 * o_bench_scale / (mult) / 4; bench_i--; ) \
      { func; func; func; func; } \
    msg(L_INFO, "%7u %s", (mult) * get_timer(&bench_timer), #func); \
  } while (0)

static void bench_libs(void)
{
  uint seed = time(NULL) + getpid();

  uint iso_state = seed;
  srand(seed);

  u64 bsd_buf[128 / 8];
  struct random_data bsd_data;
  bsd_data.state = NULL;
  if (initstate_r(seed, (char *)bsd_buf, sizeof(bsd_buf), &bsd_data) < 0)
    die("initstate_r(): %m");

  u64 bsd_buf2[128 / 8];
  errno = 0;
  initstate(seed, (char *)bsd_buf2, sizeof(bsd_buf2));
  if (errno)
    die("initstate(): %m");

  struct drand48_data svid_data;
  if (srand48_r(seed, &svid_data) < 0)
    die("srand48_r(): %m");
  srand48(seed);

  int32_t int32_res;
  long int long_res;

  msg(L_INFO, "Benchmarking libc ISO generator");
  BENCH_SLOW(100, srand(seed));
  BENCH(rand());
  BENCH(rand_r(&iso_state));

  msg(L_INFO, "Benchmarking libc BSD generator");
  BENCH_SLOW(100, srandom(seed));
  BENCH(random());
  BENCH(random_r(&bsd_data, &int32_res));

  msg(L_INFO, "Benchmarking libc SVID generator");
  BENCH(srand48(seed));
  BENCH(mrand48());
  BENCH(mrand48_r(&svid_data, &long_res));
}

static void bench_legacy(void)
{
  msg(L_INFO, "Benchmarking ucw legacy generator");
  srand(time(NULL) + getpid());
  BENCH(random_u32());
  BENCH(random_u64());
  BENCH(random_max(1000000));
  BENCH(random_max_u64(1000000));
  BENCH_SLOW(2, random_max_u64(1000000000000LL));
  BENCH_SLOW(2, random_max_u64((1LLU << 63) - 1)); // It uses slightly suboptimal intervals, so +0 is still a bad case
  BENCH_SLOW(4, random_max_u64((1LLU << 63) + 1));
  BENCH_SLOW(2, random_max(1 + random_max(1000000)));
  BENCH_SLOW(4, random_max_u64(1 + random_max_u64(1000000000000LL)));
}

static void bench_fast(void)
{
  msg(L_INFO, "Benchmarking fast random generator");
  struct fastrand *r = fastrand_new();

  if (!o_bench_only_safe)
    {
      // Can produce huge log with LOCAL_DEBUG
      uint seed = time(NULL) + getpid();
      BENCH_SLOW(100, fastrand_gen_seed(r));
      BENCH_SLOW(100, fastrand_set_seed(r, seed));
    }

  fastrand_gen_seed(r);

  BENCH(fastrand_bits(r, 1));
  BENCH(fastrand_bits(r, 32));
  BENCH(fastrand_bits_u64(r, 64));
  BENCH(fastrand_u32(r));
  BENCH(fastrand_u64(r));
  BENCH(fastrand_max(r, 1000000));
  BENCH(fastrand_max(r, (1U << 30) + 0));
  BENCH(fastrand_max(r, (1U << 30) + 1)); // Worst case for 31 bit generators (typical RAND_MAX)
  BENCH(fastrand_max(r, (1U << 31) + 0));
  BENCH(fastrand_max(r, (1U << 31) + 1)); // Worst case for 32 bit generators
  BENCH(fastrand_max_u64(r, 1000000));
  BENCH_SLOW(2, fastrand_max_u64(r, 1000000000000LL));
  BENCH_SLOW(2, fastrand_max_u64(r, (1LLU << 63) + 0));
  BENCH_SLOW(4, fastrand_max_u64(r, (1LLU << 63) + 1));
  BENCH_SLOW(2, fastrand_max(r, 1 + fastrand_max(r, 1000000)));
  BENCH_SLOW(4, fastrand_max_u64(r, 1 + fastrand_max_u64(r, 1000000000000LL)));
  BENCH_SLOW(4, fastrand_double(r));

  byte buf[1000];
  BENCH(fastrand_mem(r, buf, 1));
  BENCH(fastrand_mem(r, buf, 4));
  BENCH_SLOW(10, fastrand_mem(r, buf, 100));
  BENCH_SLOW(100, fastrand_mem(r, buf, 1000));

  fastrand_delete(r);
}

static void bench_strong(void)
{
  byte buf[100];
  struct strongrand *r;

  msg(L_INFO, "Benchmarking unbuffered strong random generator");
  r = strongrand_std_open(STRONGRAND_STD_URANDOM, 0);
  BENCH_SLOW(50, strongrand_mem(r, buf, 1));
  BENCH_SLOW(200, strongrand_mem(r, buf, 100));
  strongrand_close(r);

  msg(L_INFO, "Benchmarking buffered strong random generator");
  r = strongrand_std_open(STRONGRAND_STD_URANDOM, 128);
  BENCH_SLOW(50, strongrand_mem(r, buf, 1));
  BENCH_SLOW(200, strongrand_mem(r, buf, 100));
  strongrand_close(r);

  msg(L_INFO, "Benchmarking buffered strong random generator with autoreset");
  r = strongrand_std_open(STRONGRAND_STD_URANDOM | STRONGRAND_STD_AUTO_RESET, 128);
  BENCH_SLOW(50, strongrand_mem(r, buf, 1));
  BENCH_SLOW(200, strongrand_mem(r, buf, 100));
  strongrand_close(r);
}

int main(int argc UNUSED, char **argv)
{
  cf_def_file = INSTALL_CONFIG_DIR "/common";
  opt_parse(&options, argv + 1);

  if (o_test_fast)
    test_fast();
  if (o_test_strong)
    test_strong();

  if (o_bench_all || o_bench_libs)
    bench_libs();
  if (o_bench_all || o_bench_legacy)
    bench_legacy();
  if (o_bench_all || o_bench_fast)
    bench_fast();
  if (o_bench_all || o_bench_strong)
    bench_strong();

  return 0;
}