/**
 * Utilities related to collections.
 */
package theorycrafter.fitting.utils

import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.contract


/**
 * If the value is `null`, returns an empty collection; otherwise returns a collection containing only the value.
 */
fun <T> T?.emptyOrSingleton(): Collection<T> = listOfNotNull(this)


/**
 * Returns a new list which would is obtained by inserting the given item at the given index.
 * If the index is `null`, the item is added at the end of the list.
 */
operator fun <T> List<T>.plus(itemWithIndex: Pair<T, Int?>): List<T> {
    val (item, index) = itemWithIndex
    if ((index == null) || (index == size))
        return this + item

    val list = this
    return buildList(size + 1) {
        addAll(list.subList(0, index))
        add(item)
        addAll(list.subList(index, list.size))
    }
}


/**
 * Returns the average of the values returned by the given function on the items of this collection.
 */
inline fun <T> Collection<T>.averageOf(value: (T) -> Double): Double{
    var sum = 0.0
    var count = 0
    for (item in this) {
        sum += value(item)
        ++count
    }
    return if (count == 0) Double.NaN else sum / count
}


/**
 * Returns the suffix of the list, starting at [fromIndex].
 * Note that the returned list is a view over the original list.
 */
fun <T> List<T>.tailFrom(fromIndex: Int) = subList(fromIndex = fromIndex, toIndex = size)


/**
 * Returns the last [count] items of the list.
 * Note that the returned list is a view over the original list.
 */
fun <T> List<T>.tailOf(count: Int) = subList(fromIndex = size - count, toIndex = size)


/**
 * Associates the values in the list with their indices.
 */
fun <T> List<T>.associateWithIndex(): Map<T, Int> = buildMap(size) {
    forEachIndexed { index, value -> put(value, index) }
}


/**
 * Caches the results of the given computation.
 *
 * Returns a function that returns the same values as [computation] for the same arguments.
 * The first time it is called with a certain argument, it calls [computation] to obtain the value and caches it.
 * On subsequent calls with the same argument, it returns the cached value.
 */
fun <K, V> memoize(computation: (K) -> V): (K) -> V {
    val cache = mutableMapOf<K, V>()
    return { key: K ->
        cache.computeIfAbsent(key, computation)
    }
}


/**
 * Same as [Iterable.forEach], but doesn't create an iterator object.
 */
@OptIn(ExperimentalContracts::class)
inline fun <T> List<T>.fastForEach(action: (T) -> Unit) {
    contract { callsInPlace(action) }
    for (index in indices) {
        val item = get(index)
        action(item)
    }
}


/**
 * Like the stdlib [sumOf], but null entries are ignored.
 */
inline fun <T> Iterable<T>.sumOfNullable(selector: (T) -> Double?): Double = sumOf { selector(it) ?: 0.0 }


/**
 * Returns an [Iterable] transforming each item of the given [Iterable].
 *
 * The advantage of this function over [Iterable.map] is that it only allocates O(1) memory, for the iterator
 * object itself. The mapping occurs as the iterator is iterated over.
 */
fun <T, R> Iterable<T>.mapEach(transform: (T) -> R): Iterable<R> = TransformingIterable(this, transform)


/**
 * An [Iterable] that transforms items from another iterable.
 */
private class TransformingIterable<T, R>(
    val original: Iterable<T>,
    val transform: (T) -> R
): Iterable<R> {
    override fun iterator(): Iterator<R> = object: Iterator<R> {
        val originalIterator = original.iterator()
        override fun hasNext(): Boolean = originalIterator.hasNext()
        override fun next(): R = transform(originalIterator.next())
    }
}
