package br.com.colman.kaucasus

import io.kotest.core.spec.style.FunSpec
import io.kotest.matchers.collections.shouldHaveSize
import io.kotest.matchers.maps.shouldContainAll
import io.kotest.matchers.shouldBe

class DoubleStatisticsTest : FunSpec({

    val doubleVector = sequenceOf(0.0, 1.0, 3.0, 5.0, 11.0)
    val groups = sequenceOf("A", "B", "B", "C", "C")

    test("Sum by") {
        val r = mapOf("A" to 0.0, "B" to 4.0, "C" to 16.0)

        groups.zip(doubleVector).sumBy() shouldContainAll mapOf("A" to r["A"], "B" to r["B"], "C" to r["C"])

        groups.zip(doubleVector)
            .sumBy(keySelector = { it.first }, doubleSelector = { it.second }) shouldContainAll mapOf(
            "A" to r["A"],
            "B" to r["B"]
        )
    }

    test("Average By") {
        val r = mapOf("A" to 0.0, "B" to 2.0, "C" to 8.0)

        groups.zip(doubleVector).averageBy(
            keySelector = { it.first },
            doubleSelector = { it.second }
        ) shouldBe r
    }

    test("Bin test") {
        val binned = sequenceOf(
            doubleVector,
            doubleVector.map { it + 100.0 },
            doubleVector.map { it + 200.0 }
        ).flatMap { it }
            .zip(groups.repeat())
            .binByDouble(
                binSize = 100.0,
                valueSelector = { it.first },
                rangeStart = 0.0
            )

        binned.bins shouldHaveSize 3
        binned[5.0]!!.range.let { it.lowerBound shouldBe 0.0; it.upperBound shouldBe 100.0 }
        binned[105.0]!!.range.let { it.lowerBound shouldBe 100.0; it.upperBound shouldBe 200.0 }
        binned[205.0]!!.range.let { it.lowerBound shouldBe 200.0; it.upperBound shouldBe 300.0 }
    }
})

private fun <T> Sequence<T>.repeat(): Sequence<T> = sequence {
    while (true) yieldAll(this@repeat)
}
