/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// -*- c++ -*-

#include <faiss/utils/random.h>

extern "C" {
int sgemm_(
        const char* transa,
        const char* transb,
        FINTEGER* m,
        FINTEGER* n,
        FINTEGER* k,
        const float* alpha,
        const float* a,
        FINTEGER* lda,
        const float* b,
        FINTEGER* ldb,
        float* beta,
        float* c,
        FINTEGER* ldc);
}

namespace faiss {

/**************************************************
 * Random data generation functions
 **************************************************/

RandomGenerator::RandomGenerator(int64_t seed) : mt((unsigned int)seed) {}

int RandomGenerator::rand_int() {
    return mt() & 0x7fffffff;
}

int64_t RandomGenerator::rand_int64() {
    return int64_t(rand_int()) | int64_t(rand_int()) << 31;
}

int RandomGenerator::rand_int(int max) {
    return mt() % max;
}

float RandomGenerator::rand_float() {
    return mt() / float(mt.max());
}

double RandomGenerator::rand_double() {
    return mt() / double(mt.max());
}

SplitMix64RandomGenerator::SplitMix64RandomGenerator(int64_t seed)
        : state{static_cast<uint64_t>(seed)} {}

int SplitMix64RandomGenerator::rand_int() {
    return next() & 0x7fffffff;
}

int64_t SplitMix64RandomGenerator::rand_int64() {
    uint64_t value = next();
    return static_cast<int64_t>(value & 0x7fffffffffffffffULL);
}

int SplitMix64RandomGenerator::rand_int(int max) {
    return next() % max;
}

float SplitMix64RandomGenerator::rand_float() {
    return next() / float(std::numeric_limits<uint64_t>::max());
}

double SplitMix64RandomGenerator::rand_double() {
    return next() / double(std::numeric_limits<uint64_t>::max());
}

uint64_t SplitMix64RandomGenerator::next() {
    uint64_t z = (state += 0x9e3779b97f4a7c15ULL);
    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
    z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
    return z ^ (z >> 31);
}

/***********************************************************************
 * Random functions in this C file only exist because Torch
 *  counterparts are slow and not multi-threaded.  Typical use is for
 *  more than 1-100 billion values. */

/* Generate a set of random floating point values such that x[i] in [0,1]
   multi-threading. For this reason, we rely on re-entrant functions.  */
void float_rand(float* x, size_t n, int64_t seed) {
    // only try to parallelize on large enough arrays
    const size_t nblock = n < 1024 ? 1 : 1024;

    RandomGenerator rng0(seed);
    int a0 = rng0.rand_int(), b0 = rng0.rand_int();

#pragma omp parallel for
    for (int64_t j = 0; j < nblock; j++) {
        RandomGenerator rng(a0 + j * b0);

        const size_t istart = j * n / nblock;
        const size_t iend = (j + 1) * n / nblock;

        for (size_t i = istart; i < iend; i++) {
            x[i] = rng.rand_float();
        }
    }
}

void float_randn(float* x, size_t n, int64_t seed) {
    // only try to parallelize on large enough arrays
    const size_t nblock = n < 1024 ? 1 : 1024;

    RandomGenerator rng0(seed);
    int a0 = rng0.rand_int(), b0 = rng0.rand_int();

#pragma omp parallel for
    for (int64_t j = 0; j < nblock; j++) {
        RandomGenerator rng(a0 + j * b0);

        double a = 0, b = 0, s = 0;
        int state = 0; /* generate two number per "do-while" loop */

        const size_t istart = j * n / nblock;
        const size_t iend = (j + 1) * n / nblock;

        for (size_t i = istart; i < iend; i++) {
            /* Marsaglia's method (see Knuth) */
            if (state == 0) {
                do {
                    a = 2.0 * rng.rand_double() - 1;
                    b = 2.0 * rng.rand_double() - 1;
                    s = a * a + b * b;
                } while (s >= 1.0);
                x[i] = a * sqrt(-2.0 * log(s) / s);
            } else {
                x[i] = b * sqrt(-2.0 * log(s) / s);
            }
            state = 1 - state;
        }
    }
}

/* Integer versions */
void int64_rand(int64_t* x, size_t n, int64_t seed) {
    // only try to parallelize on large enough arrays
    const size_t nblock = n < 1024 ? 1 : 1024;

    RandomGenerator rng0(seed);
    int a0 = rng0.rand_int(), b0 = rng0.rand_int();

#pragma omp parallel for
    for (int64_t j = 0; j < nblock; j++) {
        RandomGenerator rng(a0 + j * b0);

        const size_t istart = j * n / nblock;
        const size_t iend = (j + 1) * n / nblock;
        for (size_t i = istart; i < iend; i++) {
            x[i] = rng.rand_int64();
        }
    }
}

void int64_rand_max(int64_t* x, size_t n, uint64_t max, int64_t seed) {
    // only try to parallelize on large enough arrays
    const size_t nblock = n < 1024 ? 1 : 1024;

    RandomGenerator rng0(seed);
    int a0 = rng0.rand_int(), b0 = rng0.rand_int();

#pragma omp parallel for
    for (int64_t j = 0; j < nblock; j++) {
        RandomGenerator rng(a0 + j * b0);

        const size_t istart = j * n / nblock;
        const size_t iend = (j + 1) * n / nblock;
        for (size_t i = istart; i < iend; i++) {
            x[i] = rng.rand_int64() % max;
        }
    }
}

void rand_perm(int* perm, size_t n, int64_t seed) {
    for (size_t i = 0; i < n; i++) {
        perm[i] = i;
    }

    RandomGenerator rng(seed);

    for (size_t i = 0; i + 1 < n; i++) {
        int i2 = i + rng.rand_int(n - i);
        std::swap(perm[i], perm[i2]);
    }
}

void rand_perm_splitmix64(int* perm, size_t n, int64_t seed) {
    for (size_t i = 0; i < n; i++) {
        perm[i] = i;
    }

    SplitMix64RandomGenerator rng(seed);

    for (size_t i = 0; i + 1 < n; i++) {
        int i2 = i + rng.rand_int(n - i);
        std::swap(perm[i], perm[i2]);
    }
}

void byte_rand(uint8_t* x, size_t n, int64_t seed) {
    // only try to parallelize on large enough arrays
    const size_t nblock = n < 1024 ? 1 : 1024;

    RandomGenerator rng0(seed);
    int a0 = rng0.rand_int(), b0 = rng0.rand_int();

#pragma omp parallel for
    for (int64_t j = 0; j < nblock; j++) {
        RandomGenerator rng(a0 + j * b0);

        const size_t istart = j * n / nblock;
        const size_t iend = (j + 1) * n / nblock;

        size_t i;
        for (i = istart; i < iend; i++) {
            x[i] = rng.rand_int64();
        }
    }
}

void rand_smooth_vectors(size_t n, size_t d, float* x, int64_t seed) {
    size_t d1 = 10;
    std::vector<float> x1(n * d1);
    float_randn(x1.data(), x1.size(), seed);
    std::vector<float> rot(d1 * d);
    float_rand(rot.data(), rot.size(), seed + 1);

    { //
        FINTEGER di = d, d1i = d1, ni = n;
        float one = 1.0, zero = 0.0;
        sgemm_("Not transposed",
               "Not transposed", // natural order
               &di,
               &ni,
               &d1i,
               &one,
               rot.data(),
               &di, // rotation matrix
               x1.data(),
               &d1i, // second term
               &zero,
               x,
               &di);
    }

    std::vector<float> scales(d);
    float_rand(scales.data(), d, seed + 2);

#pragma omp parallel for if (n * d > 10000)
    for (int64_t i = 0; i < n; i++) {
        for (size_t j = 0; j < d; j++) {
            x[i * d + j] = sinf(x[i * d + j] * (scales[j] * 4 + 0.1));
        }
    }
}

} // namespace faiss
