//@file:Repository("*mavenLocal")
//@file:DependsOn("com.github.holgerbrandl:kalasim:0.7-SNAPSHOT")

@file:DependsOn("com.github.holgerbrandl:kalasim:0.6.92")
@file:DependsOn("com.github.holgerbrandl:kravis:0.8.1")

import org.kalasim.*

import kravis.geomBar
import kravis.geomCol
import kravis.plot
import org.apache.commons.math3.distribution.LogNormalDistribution

import java.lang.Double.min
import kotlin.random.Random

//LogNormalDistribution(3.5, 0.88).let{ dist ->
//    repeat(100){ println(dist.sample())}
//}

//TODO set random seed, blocked by https://github.com/Kotlin/kotlin-jupyter/issues/345


class SquidGame(
    val numSteps: Int = 18,
    val numPlayers: Int = 16,
    val maxDuration: Int = 12 * 60
) : Environment(randomSeed =Random.nextInt()) {

    // randomization
    val stepTime = LogNormalDistribution(rg, 3.0, 0.88)
//    val stepTime = uniform(10,30)

    val decision = enumerated(true, false)

    // state
    var stepsLeft = numSteps
    var survivors= mutableListOf<Int>()

    val numTrials: Int
        get() = numSteps - survivors.size

    val numSurvivors : Int
       get() = survivors.size

    fun playerSurvived(playerNo: Int) = survivors.contains(playerNo)

    init {
        object : Component() {
            override fun process() = sequence {
               queue@
               for(player in 1..numPlayers){
                    hold(min(stepTime(), 100.0)) // cap time at 100sec

                    while(stepsLeft-- > 0){
                        if(decision()) continue@queue
                        hold(min(stepTime(), 100.0)) // cap time at 100sec
                    }

                    if(now > maxDuration) break

                    survivors.add(player)
                }
            }
        }
    }
}

val sim = SquidGame()
sim.run()

println("${sim.numSurvivors} survived")

sim.playerSurvived(13)

(1..18).map{ sim.playerSurvived(it)}

val manyGames = org.kalasim.misc.repeat(10000) {
    SquidGame().apply { run() }
}

val avgSurvivors = manyGames.map { it.numSurvivors }.average()
println("The average number of survivors is ${avgSurvivors}")

manyGames.plot(x = { numSurvivors }).geomBar().labs(
    title = "Outcomes of 10,000 trials",
    x = "Number of survivors",
    y = "Count"
)


val survivalProbByNo = (1..manyGames.first().numPlayers).map { playerNo ->
    playerNo to manyGames.count { it.playerSurvived(playerNo) }.toDouble() / manyGames.size
}

survivalProbByNo.plot(x = { it.first }, y = { it.second }).geomCol().labs(
    title = "Probability of survival based on player order number",
    x = "Player Order Number",
    y = "Probability"
)

val probLT2Players = manyGames.count { it.numSurvivors < 2 }.toDouble() / manyGames.size
println("the probability for less than 2 players is ${probLT2Players}")


val stepSims = (10..30).flatMap { numSteps ->
    org.kalasim.misc.repeat(10000) {
        SquidGame(numSteps = numSteps).apply { run() }
    }
}

val stepSimSummary = stepSims.groupBy { it.numSteps }.map { (steps, games) ->
    steps to games.count { it.numSurvivors < 2 }.toDouble() / games.size
}

stepSimSummary.plot(x = { it.first }, y = { it.second }).geomCol().labs(
    title = "Probability of having less than two remaining players",
    x = "Number of bridge steps",
    y = "Probability"
)