KStarbound/src/main/kotlin/ru/dbotthepony/kstarbound/util/random/RandomUtils.kt

package ru.dbotthepony.kstarbound.util.random

import com.google.gson.JsonArray
import com.google.gson.JsonElement
import it.unimi.dsi.fastutil.bytes.ByteConsumer
import org.classdump.luna.ByteString
import ru.dbotthepony.kommons.util.IStruct2d
import ru.dbotthepony.kommons.util.IStruct2i
import ru.dbotthepony.kommons.util.XXHash32
import ru.dbotthepony.kommons.util.XXHash64
import java.util.*
import java.util.random.RandomGenerator
import java.util.stream.IntStream
import kotlin.math.ceil
import kotlin.math.floor
import kotlin.math.ln
import kotlin.math.sqrt

/**
 * Decouples implementation from semantics.
 * Code should construct random generator using this function.
 * Replacing generator returned here will affect all random generation code.
 */
fun random(seed: Long = System.nanoTime()): RandomGenerator {
	return MWCRandom(seed.toULong())
}

private fun toBytes(accept: ByteConsumer, value: Short) {
	accept.accept(value.toByte())
	accept.accept((value.toInt() ushr 8).toByte())
}

private fun toBytes(accept: ByteConsumer, value: Int) {
	accept.accept(value.toByte())
	accept.accept((value ushr 8).toByte())
	accept.accept((value ushr 16).toByte())
	accept.accept((value ushr 24).toByte())
}

private fun toBytes(accept: ByteConsumer, value: Long) {
	accept.accept(value.toByte())
	accept.accept((value ushr 8).toByte())
	accept.accept((value ushr 16).toByte())
	accept.accept((value ushr 24).toByte())
	accept.accept((value ushr 32).toByte())
	accept.accept((value ushr 40).toByte())
	accept.accept((value ushr 48).toByte())
	accept.accept((value ushr 56).toByte())
}

private fun toBytes(accept: ByteConsumer, value: Double) {
	toBytes(accept, value.toBits())
}

private fun toBytes(accept: ByteConsumer, value: Float) {
	toBytes(accept, value.toBits())
}

fun staticRandom32(vararg values: Any?): Int {
	val digest = XXHash32(2938728349.toInt())

	for (value in values) {
		when (value) {
			is String -> digest.update(value.toByteArray())
			is ByteString -> digest.update(value.bytes)
			is Byte -> digest.update(value)
			is Boolean -> digest.update(if (value) 1 else 0)
			is Short -> toBytes(digest::update, value)
			is Int -> toBytes(digest::update, value)
			is Long -> toBytes(digest::update, value)
			is Double -> toBytes(digest::update, value)
			is Float -> toBytes(digest::update, value)
			null -> {} // do nothing?
			else -> throw IllegalArgumentException("Can't hash value of type ${value::class.qualifiedName}")
		}
	}

	return digest.digestAsInt()
}

fun staticRandomFloat(vararg values: Any?): Float {
	return staticRandom32(*values).ushr(8) * 5.9604645E-8f
}

fun staticRandomDouble(vararg values: Any?): Double {
	return staticRandom64(*values).ushr(11) * 1.1102230246251565E-16
}

fun staticRandomInt(min: Int, max: Int, vararg values: Any?): Int {
	val hash = staticRandomDouble(*values)
	return (min + (max - min + 1) * hash).toInt()
}

fun staticRandomLong(min: Long, max: Long, vararg values: Any?): Long {
	val hash = staticRandomDouble(*values)
	return (min + (max - min + 1L) * hash).toLong()
}

fun staticRandom64(vararg values: Any?): Long {
	val digest = XXHash64(1997293021376312589L)

	for (value in values) {
		when (value) {
			is String -> digest.update(value.toByteArray())
			is ByteString -> digest.update(value.bytes)
			is Byte -> digest.update(value)
			is Boolean -> digest.update(if (value) 1 else 0)
			is Short -> toBytes(digest::update, value)
			is Int -> toBytes(digest::update, value)
			is Long -> toBytes(digest::update, value)
			is Double -> toBytes(digest::update, value)
			is Float -> toBytes(digest::update, value)
			null -> {} // do nothing?
			else -> throw IllegalArgumentException("Can't hash value of type ${value::class.qualifiedName}")
		}
	}

	return digest.digestAsLong()
}

// normal distribution via Box-Muller
fun RandomGenerator.nextNormalFloat(stddev: Float, mean: Float): Float {
	var rand1: Float
	var rand2: Float
	var distSqr: Float

	do {
		rand1 = 2f * nextFloat() - 1f
		rand2 = 2f * nextFloat() - 1f
		distSqr = rand1 * rand1 + rand2 * rand2
	} while (distSqr >= 1)

	val mapping = sqrt(-2f * ln(distSqr) / distSqr)
	return rand1 * mapping * stddev + mean
}

// normal distribution via Box-Muller
fun RandomGenerator.nextNormalDouble(stddev: Double, mean: Double): Double {
	var rand1: Double
	var rand2: Double
	var distSqr: Double

	do {
		rand1 = 2.0 * nextDouble() - 1.0
		rand2 = 2.0 * nextDouble() - 1.0
		distSqr = rand1 * rand1 + rand2 * rand2
	} while (distSqr >= 1)

	val mapping = sqrt(-2.0 * ln(distSqr) / distSqr)
	return rand1 * mapping * stddev + mean
}

fun RandomGenerator.stochasticRound(value: Double): Long {
	val part = value - floor(value)

	if (nextDouble() < part)
		return ceil(value).toLong()
	else
		return floor(value).toLong()
}

fun RandomGenerator.bytes(): IntStream {
	// avoid flatMap since that would create lots of heap garbage
	val data = IntArray(4)
	var index = 0
	var populated = false

	return IntStream.generate {
		if (!populated || index == 4) {
			val value = nextInt()
			data[0] = value.and(0xFF)
			data[1] = value.ushr(8).and(0xFF)
			data[2] = value.ushr(16).and(0xFF)
			data[3] = value.ushr(24).and(0xFF)
			populated = true
			index = 0
		}

		data[index++]
	}
}

fun RandomGenerator.nextBytes(buf: ByteArray, offset: Int, length: Int) {
	Objects.checkFromIndexSize(offset, buf.size, length)
	if (length == 0) return
	var offset = offset

	bytes().limit(length.toLong()).forEach {
		buf[offset++] = it.toByte()
	}
}

fun RandomGenerator.nextBytes(length: Int): ByteArray {
	if (length == 0) return ByteArray(0)
	require(length > 0) { "Invalid length: $length" }
	val data = ByteArray(length)
	var index = 0

	bytes().limit(length.toLong()).forEach {
		data[index++] = it.toByte()
	}

	return data
}

fun <T> Collection<T>.random(random: RandomGenerator): T {
	if (isEmpty())
		throw NoSuchElementException("List is empty")

	return elementAt(random.nextInt(size))
}

fun JsonArray.random(random: RandomGenerator): JsonElement {
	if (isEmpty())
		throw NoSuchElementException("List is empty")

	return elementAt(random.nextInt(size()))
}

fun <T> Collection<T>.random(random: RandomGenerator, default: () -> T): T {
	if (isEmpty())
		return default.invoke()

	return elementAt(random.nextInt(size))
}

fun JsonArray.random(random: RandomGenerator, default: () -> JsonElement): JsonElement {
	if (isEmpty())
		return default.invoke()

	return elementAt(random.nextInt(size()))
}

fun RandomGenerator.nextRange(range: IStruct2i): Int {
	return if (range.component1() == range.component2()) return range.component1() else nextInt(range.component1(), range.component2())
}

fun RandomGenerator.nextRange(range: IStruct2d): Double {
	return if (range.component1() == range.component2()) return range.component1() else nextDouble(range.component1(), range.component2())
}

fun <T> MutableList<T>.shuffle(random: RandomGenerator) {
	for (i in 0 until size) {
		val rand = random.nextInt(size)
		val a = this[i]
		val b = this[rand]

		this[i] = b
		this[rand] = a
	}
}