Pith - kiln
package com.vgmlr.kiln

import java.time.LocalDate
import java.time.temporal.ChronoUnit
import kotlin.math.pow

class KilnPmsModel(
    val lead: Int,
    val active: Boolean,
    val qualified: Int,
    private val detected: Map<Long, Int>
) {
    fun leadBefore(nextStart: LocalDate): Int =
        if (active) detected[nextStart.toEpochDay()] ?: lead else lead
}

object KilnMoodInfluence {
    private val MOOD_DRAG = floatArrayOf(2f, 1f, -1f, 1f, 2f)
    private const val SYMPTOM_DRAG = 0.8f

    private const val RANGE = 3f
    private const val ELEVATION = 0.3f
    private const val ONSET = RANGE * ELEVATION / 2f

    const val MAX_LEAD = 14
    private const val PRE_DAYS = 7L
    private const val BASE_FROM = 4L
    private const val BASE_TO = 11L
    private const val CLEAR = KilnGradientPalette.GREEN_PEAK + 1
    private const val GAP = 2

    private const val MIN_LOGGED = 4
    private const val MIN_CYCLES = 2
    private const val HALF_LIFE = 4f
    private const val PRIOR_WEIGHT = 1f

    fun weight(records: List<KilnDayRecord>, buttons: KilnButtonSet): Float {
        val starts = KilnCyclePredictor.periodStarts(records, buttons)
        val cycle = KilnCyclePredictor.avgCycleLength(starts)
        val q = model(records, buttons, starts, cycle).qualified
        return (q.toFloat() / MIN_CYCLES).coerceIn(0f, 1f)
    }

    fun model(
        records: List<KilnDayRecord>,
        buttons: KilnButtonSet,
        starts: List<LocalDate>,
        cycle: Int
    ): KilnPmsModel {
        val prior = KilnGradientPalette.baseLead(cycle)
        val byDay = records.associateBy { it.epochDay }
        val runs = KilnCyclePredictor.periodRuns(records, buttons).toMap()
        val closed = KilnCyclePredictor.cycleSeries(starts)

        val samples = ArrayList<Pair<Int, Float>>()
        val detected = HashMap<Long, Int>()
        closed.forEachIndexed { i, (next, len) ->
            val start = next.minusDays(len.toLong())
            val lead = cycleLead(byDay, buttons, start, next, runs[start] ?: 1)
                ?: return@forEachIndexed
            val w = 0.5f.pow((closed.lastIndex - i) / HALF_LIFE)
            samples += (if (lead > 0) lead else prior) to w
            if (lead > 0) detected[next.toEpochDay()] = lead
        }
        val active = samples.size >= MIN_CYCLES
        val lead = if (active) weightedMedian(samples + (prior to PRIOR_WEIGHT)) else prior
        return KilnPmsModel(lead, active, samples.size, detected)
    }

    // null=skip, 0=none
    private fun cycleLead(
        byDay: Map<Long, KilnDayRecord>,
        set: KilnButtonSet,
        start: LocalDate,
        next: LocalDate,
        run: Int
    ): Int? {
        val len = ChronoUnit.DAYS.between(start, next)
        fun rec(d: LocalDate) = byDay[d.toEpochDay()]
        fun logged(d: LocalDate) = rec(d)?.let { it.mood != null || set.hasSymptom(it) } == true
        fun score(d: LocalDate): Float? = rec(d)?.takeIf { logged(d) }?.let { r ->
            (r.mood?.let { MOOD_DRAG.getOrElse(it) { 0f } } ?: 0f) +
                    SYMPTOM_DRAG * set.symptomCount(r)
        }

        val base = (maxOf(BASE_FROM, run.toLong())..BASE_TO)
            .filter { it < len - PRE_DAYS }
            .map(start::plusDays)
        val pre = (1L..PRE_DAYS).map(next::minusDays)
        if (base.count(::logged) < MIN_LOGGED || pre.count(::logged) < MIN_LOGGED) return null

        val b = base.mapNotNull(::score).average().toFloat()
        if (pre.mapNotNull(::score).average() - b < RANGE * ELEVATION) return 0

        val span = minOf(MAX_LEAD.toLong(), len - maxOf(CLEAR, run))
        if (span < 1) return 0
        val raw = (1L..span).map { score(next.minusDays(it)) }
        val thr = b + ONSET
        var lead = 0
        for (i in raw.indices) {
            val s = raw[i]?.let { raw.subList(maxOf(0, i - 1), minOf(raw.size, i + 2)) }
                ?.filterNotNull()?.average()
            if (s != null && s > thr) lead = i + 1 else if (i + 1 - lead > GAP) break
        }
        return lead
    }

    private fun weightedMedian(xs: List<Pair<Int, Float>>): Int {
        val sorted = xs.sortedBy { it.first }
        val half = sorted.sumOf { it.second.toDouble() } / 2
        var acc = 0.0
        for ((v, w) in sorted) {
            acc += w
            if (acc >= half) return v
        }
        return sorted.last().first
    }
}