package ru.dbotthepony.kstarbound.collect

import com.google.common.collect.ImmutableList
import com.google.gson.Gson
import com.google.gson.TypeAdapter
import com.google.gson.TypeAdapterFactory
import com.google.gson.reflect.TypeToken
import com.google.gson.stream.JsonReader
import com.google.gson.stream.JsonWriter
import it.unimi.dsi.fastutil.objects.ObjectArrayList
import ru.dbotthepony.kommons.util.KOptional
import java.lang.reflect.ParameterizedType
import java.util.random.RandomGenerator
import kotlin.math.absoluteValue

class WeightedList<E>(val parent: ImmutableList<Pair<Double, E>>) {
	constructor() : this(ImmutableList.of())

	val isEmpty: Boolean
		get() = parent.isEmpty()

	val isNotEmpty: Boolean
		get() = parent.isNotEmpty()

	val sum = parent.stream().mapToDouble { it.first.absoluteValue }.sum()

	fun sample(sampled: Double): KOptional<E> {
		if (isEmpty)
			return KOptional.empty()

		var sum = 0.0

		for ((v, e) in parent) {
			sum += v

			if (sum >= sampled)
				return KOptional(e)
		}

		return KOptional(parent.last().second)
	}

	fun sample(random: RandomGenerator): KOptional<E> {
		if (isEmpty)
			return KOptional.empty()

		return sample(random.nextDouble(sum))
	}

	fun sample(amount: Int, random: RandomGenerator): List<E> {
		require(amount >= 0) { "Negative amount to choose: $amount" }
		if (amount == 0 || isEmpty) return listOf()
		if (amount == parent.size) return parent.stream().map { it.second }.toList()

		// Original engine is rather lazy, because it creates a **set** of chosen indices
		// and 'while' loops until it contains desired amount of indices
		// which means the closer 'amount' is to size of weighted list, the more cpu cycles is burned

		val unseen = ObjectArrayList(parent)
		val result = ArrayList<E>()
		var currentSum = sum

		while (unseen.isNotEmpty() && result.size < amount) {
			val sampled = random.nextDouble(currentSum)
			val itr = unseen.iterator()
			var sum = 0.0

			for ((v, e) in itr) {
				sum += v

				if (sum >= sampled) {
					result.add(e)
					itr.remove()
					currentSum -= v
					break
				}
			}
		}

		return result
	}

	companion object : TypeAdapterFactory {
		override fun <T : Any> create(gson: Gson, type: TypeToken<T>): TypeAdapter<T>? {
			if (type.rawType == WeightedList::class.java) {
				val elemType = (type.type as? ParameterizedType)?.actualTypeArguments?.get(0) ?: return null

				return object : TypeAdapter<WeightedList<Any?>>() {
					private val parent = gson.getAdapter(TypeToken.getParameterized(ImmutableList::class.java, TypeToken.getParameterized(Pair::class.java, java.lang.Double::class.java, elemType).type)) as TypeAdapter<ImmutableList<Pair<Double, Any?>>>

					override fun write(out: JsonWriter, value: WeightedList<Any?>?) {
						parent.write(out, value?.parent)
					}

					override fun read(`in`: JsonReader): WeightedList<Any?>? {
						val read = parent.read(`in`) ?: return null
						return WeightedList(read)
					}
				} as TypeAdapter<T>
			}

			return null
		}
	}
}