// Copyright 2009 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include "halton.ih"
#include "rkcommon/math/vec.ih"
#include "sobol.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

///////////////////////////////////////////////////////////////////////////////
// Hash functions

inline uint32 MurmurHash3_mix(uint32 hash, uint32 k)
{
  const uint32 c1 = 0xcc9e2d51;
  const uint32 c2 = 0x1b873593;
  const uint32 r1 = 15;
  const uint32 r2 = 13;
  const uint32 m = 5;
  const uint32 n = 0xe6546b64;

  k *= c1;
  k = (k << r1) | (k >> (32 - r1));
  k *= c2;

  hash ^= k;
  hash = ((hash << r2) | (hash >> (32 - r2))) * m + n;

  return hash;
}

inline uint32 MurmurHash3_finalize(uint32 hash)
{
  hash ^= hash >> 16;
  hash *= 0x85ebca6b;
  hash ^= hash >> 13;
  hash *= 0xc2b2ae35;
  hash ^= hash >> 16;

  return hash;
}

// https://github.com/skeeto/hash-prospector/
inline uint32 tripleHash(uint32 x)
{
  x ^= x >> 17;
  x *= 0xed5ad4bb;
  x ^= x >> 11;
  x *= 0xac4c1b51;
  x ^= x >> 15;
  x *= 0x31848bab;
  x ^= x >> 14;
  return x;
}

// https://github.com/skeeto/hash-prospector/issues/12#issuecomment-1105792182
/* but https://github.com/skeeto/hash-prospector/issues/28
inline uint32 hash32(uint32 hash)
{
  hash ^= hash >> 16;
  hash *= 0x21f0aaad;
  hash ^= hash >> 15;
  hash *= 0xf35a2d97;
  hash ^= hash >> 15;

  return hash;
}*/

inline uint32 hashToRandom(uint32 value, uint32 scramble)
{
  value = (value ^ 61) ^ scramble;
  value += value << 3;
  value ^= value >> 4;
  value *= 0x27d4eb2d;
  return value;
}

///////////////////////////////////////////////////////////////////////////////
// Utility functions

inline float CranleyPattersonRotation(float x, float dx)
{
  x += dx;
  if (x >= 1.f)
    x -= 1.f;

  return x;
}

inline vec2f CranleyPattersonRotation(vec2f v, vec2f dv)
{
  const float x = CranleyPattersonRotation(v.x, dv.x);
  const float y = CranleyPattersonRotation(v.y, dv.y);

  return make_vec2f(x, y);
}

inline float radicalInverse(uint32 idx, const uint32 base)
{
  float f = 0.f, g = 1.0f, inv = 1.0f / base;

  while (idx > 0) {
    g *= inv;
    f += (idx % base) * g;
    idx /= base;
  }

  return f;
}

////////////////////V///////////////////////////////////////////////////////////
// PCG pseudo-random number generator http://www.pcg-random.org/

struct RandomSampler
{
  uint64 state;
  uint32 stream;
};

inline uint32 RandomSampler_pcg32(varying RandomSampler *uniform self)
{
  uint64 oldstate = self->state;
  self->state = oldstate * 6364136223846793005ULL + (self->stream | 1u);
  uint32 xorshifted = ((oldstate >> 18u) ^ oldstate) >> 27u;
  uint32 rot = oldstate >> 59u;
  return (xorshifted >> rot) | (xorshifted << ((-rot) & 31));
}

inline float RandomSampler_getFloat(varying RandomSampler *uniform self)
{
  return to_float_unorm(RandomSampler_pcg32(self));
}

inline void RandomSampler_init(
    varying RandomSampler *uniform self, uint32 seed, uint32 stream)
{
  self->state = 0;
  self->stream = (stream << 1u) | 1u;

  // hash seed to reduce correlation artefacts
  self->state = MurmurHash3_mix(self->state, seed);
  self->state = MurmurHash3_finalize(self->state);

  RandomSampler_pcg32(self);
  self->state += seed;
  RandomSampler_pcg32(self);
}

inline void RandomSampler_init(varying RandomSampler *uniform self, uint32 seed)
{
  RandomSampler_init(self, seed, 0);
}

///////////////////////////////////////////////////////////////////////////////
// TEA - Random numbers based on Tiny Encryption Algorithm

inline void tea8(uint32 &_v0, uint32 &_v1)
{
  uint32 v0 = _v0; // operate on registers to avoid slowdown
  uint32 v1 = _v1;
  uint32 sum = 0;

  for (uniform int i = 0; i < 8; i++) { // just 8 instead of 32 rounds
    sum += 0x9e3779b9;
    v0 += ((v1 << 4) + 0xa341316c) ^ (v1 + sum) ^ ((v1 >> 5) + 0xc8013ea4);
    v1 += ((v0 << 4) + 0xad90777d) ^ (v0 + sum) ^ ((v0 >> 5) + 0x7e95761e);
  }
  _v0 = v0;
  _v1 = v1;
}

struct RandomTEA
{
  vec2ui state;
};

inline void RandomTEA_Constructor(
    varying RandomTEA *uniform self, const uint32 idx, const uint32 seed)
{
  self->state.x = idx;
  self->state.y = seed;
}

inline varying vec2f RandomTEA_getFloats(varying RandomTEA *uniform self)
{
  tea8(self->state.x, self->state.y);
  return to_float_unorm(self->state);
}

///////////////////////////////////////////////////////////////////////////////
// LCG - Linear Congruential Generator

inline uint32 LCG_init(uint32 pixelID, uint32 sampleIndex)
{
  uint32 state = 0;
  state = MurmurHash3_mix(state, pixelID);
  state = MurmurHash3_mix(state, sampleIndex);
  state = MurmurHash3_finalize(state);
  return state;
}

inline uint32 LCG_next(uint32 value)
{
  const uint32 m = 1664525;
  const uint32 n = 1013904223;

  return value * m + n;
}

inline float LCG_getFloat(uint32 &state)
{
  state = LCG_next(state);
  return to_float_unorm(state);
}

inline vec2f LCG_getFloat2(uint32 &state)
{
  const float x = LCG_getFloat(state);
  const float y = LCG_getFloat(state);
  return make_vec2f(x, y);
}

inline vec3f LCG_getFloat3(uint32 &state)
{
  const float x = LCG_getFloat(state);
  const float y = LCG_getFloat(state);
  const float z = LCG_getFloat(state);
  return make_vec3f(x, y, z);
}

// https://psychopath.io/post/2021_01_30_building_a_better_lk_hash
inline uint32 OwenHash2(uint32 x, uint32 seed)
{
  x ^= x * 0x3d20adea;
  x += seed;
  x *= (seed >> 16) | 1;
  x ^= x * 0x05526c56;
  x ^= x * 0x53a22864;

  return x;
}

inline uint32 OwenScramble2(uint32 x, uint32 seed)
{
  x = reverseBits(x);
  x = OwenHash2(x, seed);
  return reverseBits(x);
}

inline uint32 OwenScramble4(uint32 x, uint32 seed)
{
  x = reverseBits(x);

  // https://psychopath.io/post/2022_08_14_a_fast_hash_for_base_4_owen_scrambling
  x ^= x * 0x3d20adea;
  x ^= (x >> 1) & (x << 1) & 0x55555555;
  x += seed;
  x *= (seed >> 16) | 1;
  x ^= (x >> 1) & (x << 1) & 0x55555555;
  x ^= x * 0x05526c56;
  x ^= x * 0x53a22864;

  return reverseBits(x);
}

#define SCRAMBLE0 0x9374b0d3
#define SCRAMBLE1 0x1ec38426
#define SCRAMBLE2 0xdb81c29c
#define SCRAMBLE3 0x5628372b
#define SCRAMBLE4 0xbcb745bd

#ifdef OSPRAY_PATHTRACER_DEBUG
// alternative drop-in implementations with same interface for debugging
#include "random_debug.inl"
#else

////////////////////////////////////////////////////////////////////////////////
// Low-discrepancy sampler: shuffled (padding 5 dimensions), scrambled Sobol
// Burley, "Practical Hash-based Owen Scrambling", 2020

struct LDSampler
{
  uint32 revIndex; // bit-reversed sample index
  uint32 scramble; // seed for scrambling the samples
  uint32 baseIndexLight; // when sequence is split (blue noise)
  // valid bits TODO should move to global PT Context
  uniform uint32 idxMask;
  uniform uint32 idxMaskLight;
};

inline void LDSampler_init(varying LDSampler &self,
    const uint32 seed,
    const uint32 sampleIndex,
    const uint32 sampleOffsetLight = 0,
    const uniform uint32 bits = 16,
    const uniform uint32 bitsLight = 16)
{
  self.revIndex = reverseBits(sampleIndex);
  self.scramble = seed; // avoid zero!
  self.baseIndexLight = sampleOffsetLight;
  self.idxMask = bits > 31 ? -1u : ~(-1u >> bits);
  self.idxMaskLight = bitsLight > 31 ? -1u : ~(-1u >> bitsLight);
}

inline void LDSampler_nextGroup(LDSampler &self)
{
  // padding: new decorrelated group
  self.scramble = tripleHash(self.scramble);
}

inline float LDSamplerScramble(const uint32 sample, const uint32 scramble)
{
  return to_float_unorm(reverseBits(OwenHash2(sample, scramble)));
}

inline vec2f LDSampler_get3LightSamples(const LDSampler &self,
    const uniform uint32 idx,
    const bool split,
    float &ss)
{
  uint32 revIndex =
      split ? reverseBits(self.baseIndexLight + idx) : self.revIndex;
  revIndex = OwenHash2(revIndex, self.scramble); // shuffle
  revIndex &= split ? self.idxMaskLight : self.idxMask;
  // sample
  const uint32 r0 = reverseBits(revIndex);
  const vec2ui r12 = Sobol_revSample2(revIndex);

  vec2f res;
  res.x = LDSamplerScramble(r0, self.scramble ^ SCRAMBLE0);
  res.y = LDSamplerScramble(r12.x, self.scramble ^ SCRAMBLE1);
  ss = LDSamplerScramble(r12.y, self.scramble ^ SCRAMBLE2);
  return res;
}

// the first 3 dims are already returned
// the remaining 4th dim needs finalization (scrambling)
inline vec2f LDSampler_getNext4Samples(LDSampler &self, float &ss, uint32 &u)
{
  LDSampler_nextGroup(self);
  // shuffle
  const uint32 revIndex =
      OwenHash2(self.revIndex, self.scramble) & self.idxMask;
  // sample
  const uint32 r0 = reverseBits(revIndex);
  const vec3ui r123 = Sobol_revSample3(revIndex);

  vec2f res;
  res.x = LDSamplerScramble(r0, self.scramble ^ SCRAMBLE0);
  res.y = LDSamplerScramble(r123.x, self.scramble ^ SCRAMBLE1);
  ss = LDSamplerScramble(r123.y, self.scramble ^ SCRAMBLE2);
  u = r123.z;
  return res;
}

inline float LDSampler_finalizeDim3(const LDSampler &self, const uint32 revRes)
{
  return LDSamplerScramble(revRes, self.scramble ^ SCRAMBLE3);
}

// the first 2 dims are already returned (the best ones)
// the next remaining 3 dims need finalization (scrambling)
inline vec2f LDSampler_getNext5Samples(LDSampler &self, vec3ui &ss)
{
  LDSampler_nextGroup(self);
  // shuffle
  const uint32 revIndex =
      OwenHash2(self.revIndex, self.scramble) & self.idxMask;
  // sample
  const uint32 r0 = reverseBits(revIndex);
  const vec4ui r1234 = Sobol_revSample4(revIndex);

  vec2f res;
  res.x = LDSamplerScramble(r0, self.scramble ^ SCRAMBLE0);
  res.y = LDSamplerScramble(r1234.x, self.scramble ^ SCRAMBLE1);
  ss = make_vec3ui(r1234.y, r1234.z, r1234.w);
  return res;
}

inline vec2f LDSampler_finalizeDim23(
    const LDSampler &self, const vec3ui &revRes)
{
  vec2f res;
  res.x = LDSamplerScramble(revRes.x, self.scramble ^ SCRAMBLE2);
  res.y = LDSamplerScramble(revRes.y, self.scramble ^ SCRAMBLE3);
  return res;
}

inline float LDSampler_finalizeDim4(const LDSampler &self, const vec3ui &revRes)
{
  return LDSamplerScramble(revRes.z, self.scramble ^ SCRAMBLE4);
}

#undef SCRAMBLE0
#undef SCRAMBLE1
#undef SCRAMBLE2
#undef SCRAMBLE3
#undef SCRAMBLE4

#endif

///////////////////////////////////////////////////////////////////////////////
// Halton - low discrepancy halton sequence; first two dimensions; no state

inline varying vec2f HaltonSequence_get2D(varying uint32 idx)
{
  return make_vec2f(Halton_sample2(idx), Halton_sample3(idx));
}

OSPRAY_END_ISPC_NAMESPACE