package coffee.liz.rookery.srs

import kotlin.math.exp
import kotlin.math.max
import kotlin.math.min
import kotlin.math.pow
import kotlin.random.Random

/** Answer quality, FSRS-style. */
enum class Grade(val value: Int) { AGAIN(1), HARD(2), GOOD(3), EASY(4) }

/** One answered prompt. [ms] is how long the answer took (null if not timed). */
data class Review(val item: String, val at: Long, val grade: Grade, val ms: Long? = null)

/**
 * Memory state of one item: stability (days until recall probability falls to 90%) and
 * difficulty (1 easy .. 10 hard), per FSRS-4.5 with its published default weights.
 */
data class Memory(val stability: Double, val difficulty: Double, val lastReview: Long, val reps: Int, val lapses: Int) {
    /** Probability of recalling the item at [now]. */
    fun retrievability(now: Long): Double {
        val days = max(0.0, (now - lastReview) / DAY_MS)
        return (1 + days / (9 * stability)).pow(-1.0)
    }

    companion object { const val DAY_MS = 86_400_000.0 }
}

object Fsrs {
    private val w = doubleArrayOf(
        0.4872, 1.4003, 3.7145, 13.8206, 5.1618, 1.2298, 0.8975, 0.031, 1.6474,
        0.1367, 1.0461, 2.1072, 0.0793, 0.3246, 1.587, 0.2272, 2.8755,
    )

    private fun initDifficulty(g: Grade) = (w[4] - (g.value - 3) * w[5]).coerceIn(1.0, 10.0)

    fun first(r: Review) = Memory(w[r.grade.value - 1], initDifficulty(r.grade), r.at, 1, if (r.grade == Grade.AGAIN) 1 else 0)

    fun next(m: Memory, r: Review): Memory {
        val ret = m.retrievability(r.at)
        val d = m.difficulty - w[6] * (r.grade.value - 3)
        val difficulty = (w[7] * initDifficulty(Grade.GOOD) + (1 - w[7]) * d).coerceIn(1.0, 10.0)
        val stability = if (r.grade == Grade.AGAIN) {
            min(m.stability, w[11] * difficulty.pow(-w[12]) * ((m.stability + 1).pow(w[13]) - 1) * exp(w[14] * (1 - ret)))
        } else {
            val hard = if (r.grade == Grade.HARD) w[15] else 1.0
            val easy = if (r.grade == Grade.EASY) w[16] else 1.0
            m.stability * (exp(w[8]) * (11 - difficulty) * m.stability.pow(-w[9]) * (exp(w[10] * (1 - ret)) - 1) * hard * easy + 1)
        }
        return Memory(max(0.01, stability), difficulty, r.at, m.reps + 1, m.lapses + if (r.grade == Grade.AGAIN) 1 else 0)
    }

    /** Replay an item's reviews (any order) into its current memory state. */
    fun replay(reviews: List<Review>): Memory? {
        var m: Memory? = null
        for (r in reviews.sortedBy { it.at }) m = m?.let { next(it, r) } ?: first(r)
        return m
    }
}

/**
 * Grades a timed answer: wrong → AGAIN, otherwise by speed relative to [goodMs]
 * (fast enough to use mid-solve is EASY, a long think is HARD).
 */
fun gradeAnswer(correct: Boolean, ms: Long, goodMs: Long = 2_500): Grade = when {
    !correct -> Grade.AGAIN
    ms <= goodMs / 2 -> Grade.EASY
    ms <= goodMs -> Grade.GOOD
    else -> Grade.HARD
}

/**
 * Picks what to drill next. Unlike day-based flashcards, a drill session runs continuously, so items
 * are weighted rather than scheduled: forgotten-ish (low retrievability), hard, and recently-missed
 * items come up most; never-seen items are mixed in steadily; the last few items are avoided.
 */
class DrillPicker(private val random: Random = Random.Default) {
    fun weight(m: Memory?, now: Long, recentMisses: Int): Double {
        if (m == null) return 1.5 // new: fairly likely, so the whole set gets introduced
        val forget = 1 - m.retrievability(now)
        val hard = (m.difficulty - 1) / 9
        return 0.05 + 2.0 * forget + 0.8 * hard + 1.2 * recentMisses
    }

    fun <T> pick(items: List<T>, memory: (T) -> Memory?, recentMisses: (T) -> Int, avoid: Collection<T>, now: Long): T? {
        val pool = items.filter { it !in avoid }.ifEmpty { items }
        if (pool.isEmpty()) return null
        val weights = pool.map { weight(memory(it), now, recentMisses(it)) }
        var x = random.nextDouble() * weights.sum()
        for ((i, wt) in weights.withIndex()) { x -= wt; if (x <= 0) return pool[i] }
        return pool.last()
    }
}