import com.github.doyaaaaaken.kotlincsv.dsl.csvWriter
import java.io.File
import javax.sound.sampled.AudioSystem
import javax.sound.sampled.UnsupportedAudioFileException
import kotlin.math.log2
import kotlin.random.Random
import kotlin.system.measureNanoTime
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import org.tritonus.share.sampled.file.TAudioFileFormat

private fun <T> benchmark(block: () -> T): T {
    val result: T
    val time = measureNanoTime {
        result = block()
    }
    println("Took ${time / 1000000.0} ms")

    return result
}

fun main() {
//    entropyTesting()
    entropySweepNTarget1H()
    entropySweepNTarget2H()
    entropySweep20NTarget1H()
    entropySweepTargetN100()
    entropySweepTargetN500()
}

// interested in:
// entropy for sweeping item size items compared to entropy for full randomness for a high target
// entropy for sweeping target length with fixed size

// Entropy could be measured in two different ways:
// same set chosen every time (this is what we are doing)
// same song chosen across times

private fun Iterable<Pair<Int, Pair<Double, Double>>>.write(name: String) {
    csvWriter().open(name) {
        writeRow(listOf("N", "RGLI", "Random"))
        forEach { row ->
            writeRow(row.first, row.second.first, row.second.second)
        }
    }
}

private fun entropySweepNTarget1H() {
    val result = sweepItems(60 * 60.0)

    result.write("sweepNTarget1H.csv")
}

private fun entropySweepNTarget2H() {
    val result = sweepItems(60 * 60.0 * 2)

    result.write("sweepNTarget2H.csv")
}

private fun entropySweep20NTarget1H() {
    val result = sweepItems(60 * 60.0, 10..30 step 1)

    result.write("sweep20NTarget1H.csv")
}

private fun entropySweepTargetN20() {
    val result = sweepTarget(20, 2000..6000 step 100)

    result.write("sweepTargetN20.csv")
}

private fun entropySweepTargetN100() {
    val result = sweepTarget(100, 22000..25000 step 100)

    result.write("sweepTargetN100.csv")
}

private fun entropySweepTargetN500() {
    val result = sweepTarget(500)

    result.write("entropySweepTargetN500.csv")
}


private fun entropyTesting() {
    val rand = Random(38462894653)
    val items = 200
    val musics = getMusic(items, rand).map(Music::toItem)

    val (rgliS, randS) = benchmark {
        computeEntropyFor(musics, 3600.0, 50000)
    }
    println("RGLI entropy: $rgliS bits")
    println("Random entropy: $randS bits")
}

private fun sweepItems(target: Double, range: IntProgression = (10..500 step 10), trials: Int = 100) = range.map { num ->
    println("On step $num")
    val items = getMusic(num).map(Music::toItem)
    num to computeEntropyFor(items, target, trials)
}

private fun sweepTarget(itemsCount: Int, range: IntProgression = (60..3600 step 100), trials: Int = 100) = range.map { num ->
    println("On step $num")
    val items = getMusic(itemsCount).map(Music::toItem)
    num to computeEntropyFor(items, num.toDouble(), trials)
}

private fun computeEntropyFor(
    items: List<Item>,
    target: Double,
    trials: Int,
): Pair<Double, Double> {
    val rgliSSP = RGLISSP()

    val (rgli, random) = runBlocking {
        val rgliJob = async {
            withContext(Dispatchers.Default) {
                (1..trials).map {
                    rgliSSP.solveSSP(items, target).items.toSet()
                }.entropyEven()
            }
        }
        val randomJob = async {
            withContext(Dispatchers.Default) {
                (1..trials).map {
                    RandomSSP.solveSSP(items.shuffled(), target).items.toSet()
                }.entropyEven()
            }
        }

        val one = rgliJob.await()
        val two = randomJob.await()
        one to two
    }

    return rgli to random
}

data class Music(
    val name: String,
    val length: Double,
    val path: String,
) {
    fun toItem() = Item(this, length)
}

private fun getMusic(count: Int, random: Random = Random.Default): List<Music> {
    return File("~/Music")
        .walkTopDown()
        .filter {
            it.isFile && it.extension == "mp3"
        }
        .map { file ->
            val fileFormat = AudioSystem.getAudioFileFormat(file)
            val duration = if (fileFormat is TAudioFileFormat) {
                val properties: Map<*, *> = fileFormat.properties()
                val key = "duration"
                val microseconds = properties[key] as Long?
                (microseconds!! / 1000.0 / 1000.0)
            } else {
                throw UnsupportedAudioFileException()
            }

            Music(
                file.name,
                duration,
                file.path,
            )
        }
        .filter {
            it.length in (1.0 * 60)..(20.0 * 60) && random.nextDouble() < 0.1
        }
        .take(count)
        .toList()
}

private fun <T> Iterable<T>.entropy(probabilityDistribution: (T) -> Double): Double =
    toSet().sumOf { item ->
        val prob = probabilityDistribution(item)
        -prob * log2(prob)
    }

fun <T> Collection<T>.entropyEven() = entropy { item ->
    val total = count { item == it }
    total.toDouble() / size.toDouble()
}

private fun distributionTest() = runBlocking {
    val avgSongLength = 60.0 * 3.0
    val songStandardDeviation = 60.0
    val rand = java.util.Random()
    var id = 0
    val nums = buildList {
        repeat(500) {
            add(
                Item(
                    id++,
                    rand.nextGaussian(avgSongLength, songStandardDeviation).coerceAtLeast(60.0)
                )
            )
        }
    }

    val target = 60.0 * 60.0 * 24.0

    val result = benchmark { RGLISSP().solveSSP(nums, target) }

    val resultSum = result.sum
    println("Result Sum: $resultSum")
    println("off by ${target - resultSum} seconds")
    println("elements: $result")
}

data class Solution(
    val items: List<Item> = emptyList(),
    val sum: Double = items.sumOf(Item::length),
) {
    fun add(item: Item) = Solution(items + item, sum + item.length)
    fun isCloseEnough(delta: Double) = delta < 0.25
    fun replace(old: Item, new: Item) = Solution(items.map { if (it == old) new else it })
}

data class Item(
    val id: Any,
    val length: Double,
)

fun interface SSPAlgorithm {
    fun solveSSP(items: List<Item>, target: Double): Solution
}

object RandomSSP : SSPAlgorithm {
    override fun solveSSP(items: List<Item>, target: Double): Solution {
        return items.fold(Solution()) { solution, item ->
            val new = solution.add(item)
            if (new.sum <= target) new else solution
        }
    }
}

class RGLISSP(val trials: Int = 50) : SSPAlgorithm {
    override fun solveSSP(items: List<Item>, target: Double): Solution {
        var best = Solution()
        repeat(trials) {
            val shuffled = items.shuffled()

            val initial = shuffled.fold(Solution()) { acc, item ->
                val new = acc.add(item)
                if (new.sum <= target) new else acc
            }

            val localImprovement = initial.items.shuffled().fold(initial) { solution, item ->
                val delta = target - solution.sum
                if (solution.isCloseEnough(delta)) return@fold solution

                val replacement = (items - solution.items.toSet()).filter {
                    it.length > item.length && it.length - item.length <= delta
                }.maxByOrNull(Item::length) ?: return@fold solution

                solution.replace(item, replacement)
            }

            if (localImprovement.sum > best.sum) {
                best = localImprovement
            }
            if (best.sum == 0.0) {
                return best
            }
        }

        return best
    }
}