import libc;
import math.(sqrt,cos,log,M_PI);
import printer.formatter.(hex, repr);

// Implementation of XOR-shift RNGs by George Marsaglia
record XORShiftRNG(
  w : UInt32,
  x : UInt32,
  y : UInt32,
  z : UInt32
);

overload XORShiftRNG() = XORShiftRNG(pr(), pr(), pr(), pr());

define generate;

[I | Integer?(I) and TypeSize(I) <= 32] overload generate(r, static I)
    = wrapCast(I, generate(r, UInt32));

overload generate(r : XORShiftRNG, static UInt32){
  var t = xsl(r.x, 15);
  r.x = r.y;
  r.y = r.z;
  r.z = r.w;
  r.w = xor(xsr(r.w, 21), xsr(t, 4));

  return r.w;
}

overload generate(r, static UInt64){
  var x = UInt64(generate(r, UInt32));
  var y = UInt64(generate(r, UInt32));

  return bitor(wrapBitshl(x, 32), y);
}

overload generate(r, static Int64) = wrapCast(Int64, generate(r, UInt64));

[n, T] overload generate(r, static Array[T, n]){
  var result = Array[T, n]();

  for(i in range(n)) result[i] <-- generate(r, i);

  return move(result);
}

define uniform;

// TODO: This is a bit of a rubbish default implementation
// It works fine for integers (I think) but is almost certainly wrong for floats
// Either way it could be a lot faster.
[T | Integer?(T)] overload uniform(rnd, from : T, to : T){
  if(safeFromOverflow?(from, to)){
    var ub = to - from;
    var mask = toMask(ub);

    var result = bitand(mask, generate(rnd, T));

    while(result >= ub) result = bitand(mask, generate(rnd, T));

    return result + from;
  } else {
    var result = generate(rnd, T);

    while(result < from or result >= to) result = generate(rnd, T);

    return result;
  }
}

overload uniform(rnd, from : Float32, to : Float32)
    = from + (to - from) * smallFloat(rnd);

overload uniform(rnd, from : Float64, to : Float64)
    = from + (to - from) * smallDouble(rnd);

overload uniform(rnd, static Float32) = uniform(rnd, 0. ,1.);
overload uniform(rnd, static Float64) = uniform(rnd, 0. ,1.);

define normal;

[T | Numeric?(T)] overload normal(rnd, mean:T, stdev:T)
    = mean + stdev * sqrt(-T(2) * log(uniform(rnd,T(0),T(1)))) * cos(T(2) * M_PI * uniform(rnd,T(0),T(1)));
[T | Numeric?(T)] overload normal(rnd, mean:T) = normal(rnd, mean, T(1));

private smallDouble(r) =
  bitcast(Double, bitor(mask(generate(r, Int64), 52), bitcast(UInt64,1.0)))-1;

private smallFloat(r) = mask(generate(r, UInt32), 24) / Float32(bitshl(1, 24));

private xor(x, y) = bitxor(x, y);
private xsl(x, n) = xor(x, wrapBitshl(x, n));
private xsr(x, n) = xor(x, bitshr(x, n));
[I] private mask(x:I, n) = bitand(x, bitshl(I(1), n) - I(1));
private pr() = UInt32(libc.random());
private bitcast(T, x) = Pointer[T](&x)^;

private define safeFromOverflow?;

[T | UnsignedInteger?(T)] overload safeFromOverflow?(x : T, y : T) = true;
[T | SignedInteger?(T)] overload safeFromOverflow?(x : T, y : T) = (x < T(0) and T(0) < y) or (y - x > T(0));

private define toMask;

private inline pb(x, n){ x = bitor(x, bitshr(x, n)); }
[T] overload toMask(x:T){
  var t = x;
  pb(t, 1);
  pb(t, 2);
  pb(t, 4);
  if (TypeSize(T) > 1) {
    pb(t, 8);
    if (TypeSize(T) > 2) {
      pb(t, 16);
      if (TypeSize(T) > 4) {
        pb(t, 32);
      }
    }
  }
  return t;
}
