#include "Random.h"
#include "Func.h"
#include "IRMutator.h"
#include "IROperator.h"

namespace Halide {
namespace Internal {

using std::string;
using std::vector;

namespace {

// Some randomly-generated integers.
#define C0 576942909
#define C1 1121052041
#define C2 1040796640

// Permute a 32-bit unsigned integer using a fixed psuedorandom
// permutation.
Expr rng32(const Expr &x) {
 internal_assert(x.type() == UInt(32));

 // A polynomial P with coefficients C0 .. CN induces a permutation
 // modulo 2^d iff:
 // 1) P(0) != P(1) modulo 2
 // 2) sum(i * Ci) is odd

 // (See http://en.wikipedia.org/wiki/Permutation_polynomial#Rings_Z.2FpkZ)

 // For a quadratic, this is only satisfied by:
 // C0 anything
 // C1 odd
 // C2 even

 // The coefficients defined above were chosen to satisfy this
 // property.

 // It's pretty random, but note that the quadratic term disappears
 // if inputs are the multiples of 2^16, and so you get a linear
 // sequence. However, *that* linear sequence probably varies in
 // the low bits, so if you run in through the permutation again,
 // you should break it up. All actual use of this runs it through
 // multiple times in order to combine several inputs, so it should
 // be ok. The other flaw is it's a permutation, so you get no
 // collisions. Birthday paradox be damned.

 // However, it's exceedingly cheap to compute, as it only uses
 // vectorizable int32 muls and adds, and the resulting numbers:
 // - Have the correct moments for a uniform distribution
 // - Have no serial correlations in any of the bits
 // - Have a completely flat power spectrum
 // - Have no visible patterns

 // So I declare this good enough for image processing.

 // If it's just a const (which it often is), save the simplifier some work:
 if (const uint64_t *i = as_const_uint(x)) {
 return make_const(UInt(32), ((C2 * (*i)) + C1) * (*i) + C0);
 }

 return (((C2 * x) + C1) * x) + C0;
}
} // namespace

Expr random_int(const vector<Expr> &e) {
 internal_assert(!e.empty());
 internal_assert(e[0].type() == Int(32) || e[0].type() == UInt(32));
 // Permute the first term
 Expr result = rng32(cast(UInt(32), e[0]));
 for (size_t i = 1; i < e.size(); i++) {
 internal_assert(e[i].type() == Int(32) || e[i].type() == UInt(32));
 // Add in the next term and permute again
 string name = unique_name('R');
 // If it's a const, save the simplifier some work
 const uint64_t *ir = as_const_uint(result);
 const uint64_t *ie = as_const_uint(e[i]);
 if (ir && ie) {
 result = rng32(make_const(UInt(32), (*ir) + (*ie)));
 } else {
 result = Let::make(name, result + cast<uint32_t>(e[i]),
 rng32(Variable::make(UInt(32), name)));
 }
 }
 // The low bytes of this have a poor period, so mix in the high bytes for
 // two additional instructions.
 result = result ^ (result >> 16);

 return result;
}

Expr random_float(const vector<Expr> &e) {
 Expr result = random_int(e);
 // Set the exponent to one, and fill the mantissa with 23 random bits.
 result = (127 << 23) | (cast<uint32_t>(result) >> 9);
 // The clamp is purely for the benefit of bounds inference.
 return clamp(reinterpret(Float(32), result) - 1.0f, 0.0f, 1.0f);
}

namespace {

class LowerRandom : public IRMutator {
 using IRMutator::visit;

 Expr visit(const Call *op) override {
 if (op->is_intrinsic(Call::random)) {
 vector<Expr> args = op->args;
 // Insert the free vars in reverse, so innermost vars typically end
 // up last.
 args.insert(args.end(), extra_args.rbegin(), extra_args.rend());
 if (op->type == Float(32)) {
 return random_float(args);
 } else if (op->type == Int(32)) {
 return cast<int32_t>(random_int(args));
 } else if (op->type == UInt(32)) {
 return random_int(args);
 } else {
 internal_error << "The intrinsic random() returns an Int(32), UInt(32) or a Float(32).\n";
 return Expr();
 }
 } else {
 return IRMutator::visit(op);
 }
 }

 vector<Expr> extra_args;

public:
 LowerRandom(const vector<VarOrRVar> &free_vars, int tag) {
 for (const VarOrRVar &v : free_vars) {
 if (v.is_rvar) {
 extra_args.push_back(v.rvar);
 } else {
 extra_args.push_back(v.var);
 }
 }
 extra_args.emplace_back(tag);
 }
};

} // namespace

Expr lower_random(const Expr &e, const vector<VarOrRVar> &free_vars, int tag) {
 LowerRandom r(free_vars, tag);
 return r.mutate(e);
}

} // namespace Internal
} // namespace Halide