package dev.privacyllc.period.domain.prediction

import java.time.LocalDate
import kotlin.math.abs
import kotlin.math.exp
import kotlin.math.ln
import kotlin.math.roundToLong

/**
 * The engine PRODUCT_PLAN.md §12 specifies.
 *
 * It keeps a **discrete probability distribution over candidate start dates**
 * (§12 step 7) rather than a date and a margin bolted onto it. Everything the
 * product needs falls out of that one structure for the same reason at the same
 * time: the most likely date is its mode, the window is the narrowest span
 * holding most of its mass, and a "Not yet" is the distribution being
 * conditioned on what the user just told us — which is what §13 asks for and
 * what a date-plus-margin design cannot express.
 *
 * ## The five decisions, and why each is that way
 *
 * **Recency weighting** (§12 step 3) — a cycle from two years ago should not
 * count like last month's. `IntervalAnalysis` applies the decay.
 *
 * **A weighted median, not a mean** (§12 step 4) — the §51 outlier case is a
 * 45-day gap in an otherwise 29-day history. A mean moves to 31.7 and the
 * forecast is wrong for a year; a weighted median does not move at all.
 *
 * **A Laplace-shaped distribution, not a normal one** — cycles have heavier
 * tails than a bell curve. Under a normal distribution a period arriving four
 * days late is nearly impossible, so the model refuses to widen and keeps being
 * confidently wrong. Laplace assigns that tail real mass, which is the honest
 * shape for this data and produces windows that contain the answer more often
 * without being uselessly wide.
 *
 * **Variability from both spread and past error** (§12 step 5) — a user whose
 * cycles agree but whose forecasts have been off by three days does not have a
 * tight window, whatever her intervals say. Ignoring measured error while
 * claiming to learn would be the app not listening to itself.
 *
 * **Trend, damped** (§12 step 6) — a consistent drift is followed, at a fraction
 * of its size. §12: "Do not overfit one cycle."
 *
 * ## What it will not do
 *
 * Show percentages. The distribution stays internal — §12 is explicit that the
 * normal interface should not surface them without usability research behind it.
 */
class PersonalPredictionEngine : PredictionEngine {

    override val modelVersion: String = "personal-1"

    override fun predict(input: PredictionInput): Prediction? {
        val starts = input.confirmedStarts.distinct().sorted()
        if (starts.isEmpty()) return null

        val lastStart = starts.last()
        val intervals = IntervalAnalysis.intervals(starts)

        val centre = centreOf(intervals)
        val scale = scaleOf(intervals, centre, input.recentAbsoluteErrors)

        // The earliest date still possible. A period that had started would have
        // been logged, so today is the floor even with no "Not yet" at all; an
        // explicit "Not yet" moves it further.
        val ruledOutThrough = maxOf(
            input.today.minusDays(1),
            input.notYet.maxOfOrNull { it.date } ?: LocalDate.MIN,
        )

        val distribution = distribution(
            centreDate = lastStart.plusDays(centre.roundToLong()),
            scale = scale,
            ruledOutThrough = ruledOutThrough,
        ) ?: return null

        val mode = distribution.maxBy { it.value }.key
        val (windowStart, windowEnd) = window(distribution, mode)

        val confidence = confidence(
            intervals = intervals,
            scale = scale,
            notYetCount = input.notYet.size,
            recentErrors = input.recentAbsoluteErrors,
        )

        return Prediction(
            mostLikelyStartDate = mode,
            windowStart = windowStart,
            windowEnd = windowEnd,
            confidenceScore = confidence,
            confidenceLabel = label(confidence, intervals.size),
            modelVersion = modelVersion,
        )
    }

    // -----------------------------------------------------------------------
    // Centre
    // -----------------------------------------------------------------------

    /**
     * The personalised centre, with recency trusted in proportion to how much
     * the user's cycles actually agree.
     *
     * ## Why recency is not weighted equally for everyone
     *
     * Recency weighting rests on an assumption: that the recent past predicts
     * the near future. For a woman whose cycles run 28, 29, 28, 29 that is
     * obviously true. For one running 25, 34, 29, 37, 26, 32 it is obviously
     * false — her most recent cycle is a draw from a wide distribution, not a
     * signal, and leaning on it means chasing noise and landing further from the
     * answer than simply taking her whole history would.
     *
     * Measured, not assumed: weighting recency equally for both cost three days
     * of accuracy on the §51 variable-user fixture while gaining nothing on the
     * stable ones. So consistency decides how much recency is trusted — full
     * weight when her cycles agree, sliding toward her whole history when they
     * do not.
     *
     * The questionable down-weighting from `IntervalAnalysis` applies either
     * way. Doubt about a data point is not the same as doubt about recency.
     */
    private fun centreOf(intervals: List<Interval>): Double {
        if (intervals.isEmpty()) return POPULATION_DEFAULT_DAYS

        // Weights carrying the questionable damping but no recency decay.
        val flat = intervals.map { interval ->
            interval.days.toDouble() to if (interval.isQuestionable) {
                IntervalAnalysis.QUESTIONABLE_WEIGHT_FACTOR
            } else {
                1.0
            }
        }

        val wholeHistory = IntervalAnalysis.weightedMedian(flat)
        val recencyLed = IntervalAnalysis.weightedMedian(intervals.map { it.days.toDouble() to it.weight })

        val totalFlat = flat.sumOf { it.second }
        val rawSpread = if (totalFlat <= 0.0) 0.0 else {
            flat.sumOf { (d, w) -> abs(d - wholeHistory) * w } / totalFlat
        }
        val consistency = 1.0 / (1.0 + rawSpread / RECENCY_TRUST_SENSITIVITY)

        val base = consistency * recencyLed + (1.0 - consistency) * wholeHistory
        return base + dampedTrend(intervals, base, rawSpread)
    }

    /**
     * A consistent drift, followed at a fraction of its size.
     *
     * Needs enough history to tell a trend from noise, and compares the recent
     * half against the older half using medians so one unusual cycle cannot
     * manufacture a trend. The result is damped hard: §12 says adapt gradually
     * and do not overfit one cycle, and a model that chases every wobble
     * produces a forecast that moves for no reason the user can see.
     */
    private fun dampedTrend(intervals: List<Interval>, base: Double, spread: Double): Double {
        if (intervals.size < MINIMUM_FOR_TREND) return 0.0

        val ordered = intervals.sortedBy { it.ageIndex }          // newest first
        val half = ordered.size / 2
        val recent = ordered.take(half).filterNot { it.isQuestionable }
        val older = ordered.drop(half).filterNot { it.isQuestionable }
        if (recent.isEmpty() || older.isEmpty()) return 0.0

        val shift = IntervalAnalysis.median(recent.map { it.days.toDouble() }) -
            IntervalAnalysis.median(older.map { it.days.toDouble() })

        // Ignore drift smaller than the noise it would be indistinguishable
        // from — and "noise" is this user's own spread, not a constant.
        //
        // A fixed one-day floor fired on a 42-day-cycle history whose medians
        // differed by a single day, turning an exact forecast into a wrong one.
        // One day is a real trend at a cycle of 28 and rounding error at 42.
        val floor = maxOf(MINIMUM_TREND_DAYS, spread * TREND_SIGNAL_RATIO)
        if (abs(shift) <= floor) return 0.0

        val damped = shift * TREND_DAMPING
        return damped.coerceIn(-base * MAX_TREND_FRACTION, base * MAX_TREND_FRACTION)
    }

    // -----------------------------------------------------------------------
    // Spread
    // -----------------------------------------------------------------------

    /**
     * How wide the distribution is, from the user's own variability and from
     * how wrong the app has recently been.
     *
     * The larger of the two rather than a blend: they are both lower bounds on
     * uncertainty, and averaging them lets good-looking intervals hide a run of
     * bad forecasts. Being wrong is evidence, and it outranks looking tidy.
     */
    private fun scaleOf(intervals: List<Interval>, centre: Double, recentErrors: List<Int>): Double {
        val fromSpread = when {
            intervals.isEmpty() -> NO_HISTORY_SCALE
            intervals.size == 1 -> SINGLE_INTERVAL_SCALE
            else -> {
                val deviations = intervals.map { abs(it.days - centre) to it.weight }

                // Two spread estimates, and the larger wins.
                //
                // The median absolute deviation is robust and, on its own, badly
                // wrong for this data: a user alternating 25 and 37 has HALF her
                // deviations at zero, so the MAD is zero and the model reads a
                // wildly unpredictable cycle as perfectly consistent. A test
                // caught it asserting §15's rule that volume alone must not buy
                // High confidence — the twenty-cycle alternating case came back
                // High.
                //
                // The mean absolute deviation sees that bimodality. It is less
                // robust to a single freak cycle, which is fine here because
                // robustness comes from somewhere better: IntervalAnalysis has
                // already down-weighted anything questionable, so an outlier
                // arrives at a quarter of its size rather than being trimmed
                // away by a statistic that cannot tell it apart from a pattern.
                val mad = IntervalAnalysis.weightedMedian(deviations) * MAD_TO_SCALE
                val totalWeight = intervals.sumOf { it.weight }
                val meanAd = if (totalWeight <= 0.0) 0.0 else {
                    deviations.sumOf { (d, w) -> d * w } / totalWeight * MEAN_AD_TO_SCALE
                }
                maxOf(mad, meanAd)
            }
        }

        val fromError = if (recentErrors.isEmpty()) 0.0 else {
            val recent = recentErrors.take(ERROR_WINDOW)
            IntervalAnalysis.median(recent.map { it.toDouble() }) * ERROR_TO_SCALE
        }

        return maxOf(fromSpread, fromError).coerceIn(MINIMUM_SCALE, MAXIMUM_SCALE)
    }

    // -----------------------------------------------------------------------
    // The distribution
    // -----------------------------------------------------------------------

    /**
     * A discretised Laplace over candidate dates, with everything ruled out
     * removed and the remainder renormalised.
     *
     * Censoring then renormalising is the whole of §13: mass does not shift by a
     * day, it is redistributed over what is still possible, so the mode, the
     * window and the confidence all move for the same reason.
     */
    private fun distribution(
        centreDate: LocalDate,
        scale: Double,
        ruledOutThrough: LocalDate,
    ): Map<LocalDate, Double>? {
        val span = (scale * SPAN_IN_SCALES).roundToLong().coerceIn(MIN_SPAN_DAYS, MAX_SPAN_DAYS)

        var candidates = (-span..span).map { offset ->
            val date = centreDate.plusDays(offset)
            date to exp(-abs(offset) / scale)
        }.filter { it.first > ruledOutThrough }

        // Every candidate ruled out — a user long past the whole window. Rather
        // than dividing by zero or asserting a date with no evidence, restart
        // the distribution from the first day still possible and say so through
        // a wider scale. §13: repeated "Not yet" keeps updating the forecast.
        if (candidates.isEmpty()) {
            val restart = ruledOutThrough.plusDays(1)
            val widened = (scale * OVERDUE_WIDENING).coerceAtMost(MAXIMUM_SCALE)
            candidates = (0..span).map { offset ->
                restart.plusDays(offset) to exp(-offset / widened)
            }
        }

        val total = candidates.sumOf { it.second }
        if (total <= 0.0) return null
        return candidates.associate { (date, mass) -> date to mass / total }
    }

    /** The narrowest run of dates around the mode holding [WINDOW_MASS] of the probability. */
    private fun window(distribution: Map<LocalDate, Double>, mode: LocalDate): Pair<LocalDate, LocalDate> {
        val dates = distribution.keys.sorted()
        val modeIndex = dates.indexOf(mode)
        var low = modeIndex
        var high = modeIndex
        var mass = distribution.getValue(mode)

        while (mass < WINDOW_MASS && (low > 0 || high < dates.lastIndex)) {
            val nextLow = if (low > 0) distribution.getValue(dates[low - 1]) else -1.0
            val nextHigh = if (high < dates.lastIndex) distribution.getValue(dates[high + 1]) else -1.0
            // Always take the heavier side, which is what makes the window the
            // narrowest one containing that mass rather than merely a symmetric
            // one — and asymmetry is the point after a "Not yet".
            if (nextHigh >= nextLow) { high++; mass += nextHigh } else { low--; mass += nextLow }
        }
        return dates[low] to dates[high]
    }

    // -----------------------------------------------------------------------
    // Confidence
    // -----------------------------------------------------------------------

    /**
     * §15's list, with agreement dominating.
     *
     * The rule that decides every argument here: **do not assign High purely
     * because the user has entered a large number of cycles.** Somebody with
     * 25, 33, 28, 37, 26, 32 has plenty of data and an unpredictable cycle, and
     * telling her the forecast is High is the failure. So evidence and accuracy
     * can only cap what agreement has already allowed.
     */
    private fun confidence(
        intervals: List<Interval>,
        scale: Double,
        notYetCount: Int,
        recentErrors: List<Int>,
    ): Double {
        if (intervals.isEmpty()) return NO_HISTORY_CONFIDENCE

        val agreement = 1.0 / (1.0 + scale / AGREEMENT_SENSITIVITY)
        val evidence = (intervals.size.toDouble() / SATURATION_CYCLES).coerceAtMost(1.0)

        val accuracy = if (recentErrors.isEmpty()) NEUTRAL_ACCURACY else {
            val mean = recentErrors.take(ERROR_WINDOW).average()
            (1.0 / (1.0 + mean / ACCURACY_SENSITIVITY)).coerceIn(0.0, 1.0)
        }

        // Questionable intervals are doubt about the data itself, which is a
        // different thing from the data disagreeing — §15 lists both.
        val questionablePenalty = intervals.count { it.isQuestionable }.toDouble() /
            intervals.size * QUESTIONABLE_PENALTY

        val notYetPenalty = (1.0 - exp(-notYetCount * NOT_YET_SENSITIVITY)) * NOT_YET_MAX_PENALTY

        return (agreement * evidence * accuracy - questionablePenalty - notYetPenalty)
            .coerceIn(0.0, 1.0)
    }

    private fun label(confidence: Double, intervalCount: Int): ConfidenceLabel = when {
        intervalCount < MINIMUM_FOR_MEDIUM -> ConfidenceLabel.LOW
        confidence >= HIGH_THRESHOLD -> ConfidenceLabel.HIGH
        confidence >= MEDIUM_THRESHOLD -> ConfidenceLabel.MEDIUM
        else -> ConfidenceLabel.LOW
    }

    private companion object {
        const val POPULATION_DEFAULT_DAYS = 28.0

        // Spread. Every one of these was tuned against the §51 fixtures rather
        // than chosen — §12: "Tune with tests rather than guessing."
        const val NO_HISTORY_SCALE = 3.5
        const val SINGLE_INTERVAL_SCALE = 2.5
        const val MAD_TO_SCALE = 1.1
        const val MEAN_AD_TO_SCALE = 1.0
        /** Spread at which recency is trusted about half as much. Tuned against the §51 fixtures. */
        const val RECENCY_TRUST_SENSITIVITY = 2.0
        const val ERROR_TO_SCALE = 0.9
        const val ERROR_WINDOW = 6
        const val MINIMUM_SCALE = 0.7
        const val MAXIMUM_SCALE = 9.0

        // Distribution shape.
        const val SPAN_IN_SCALES = 6.0
        const val MIN_SPAN_DAYS = 4L
        const val MAX_SPAN_DAYS = 45L
        const val WINDOW_MASS = 0.80
        const val OVERDUE_WIDENING = 1.5

        // Trend.
        const val MINIMUM_FOR_TREND = 4
        const val MINIMUM_TREND_DAYS = 1.0
        const val TREND_SIGNAL_RATIO = 1.0
        // Raised from 0.35 after measurement: the §51 drift fixtures were the
        // two the engine won, and it was leaving accuracy on the table by
        // following a real signal too timidly. Still well under 1.0, because
        // §12's "do not overfit one cycle" is the other half of the rule.
        const val TREND_DAMPING = 0.50
        const val MAX_TREND_FRACTION = 0.15

        // Confidence.
        const val NO_HISTORY_CONFIDENCE = 0.10
        const val AGREEMENT_SENSITIVITY = 1.6
        const val SATURATION_CYCLES = 5.0
        const val NEUTRAL_ACCURACY = 0.92
        const val ACCURACY_SENSITIVITY = 2.5
        const val QUESTIONABLE_PENALTY = 0.20
        const val NOT_YET_SENSITIVITY = 0.55
        const val NOT_YET_MAX_PENALTY = 0.30
        const val MINIMUM_FOR_MEDIUM = 2
        const val HIGH_THRESHOLD = 0.55
        const val MEDIUM_THRESHOLD = 0.30
    }
}
