/**
 * @file PoseLandmarkSmoother.kt
 * @brief Per-frame exponential-moving-average smoothing of ML Kit pose landmarks.
 */
package com.example.jnicpp.bowling

import com.google.mlkit.vision.pose.Pose

/**
 * @brief A single pose landmark's position and detection confidence,
 *        smoothed across frames by [PoseLandmarkSmoother].
 *
 * Deliberately independent of ML Kit's own `PoseLandmark`/`PointF3D` so
 * downstream consumers (angle math, skeleton drawing) don't need any ML Kit
 * types.
 *
 * @param x Smoothed x coordinate, in analysis-image pixel space.
 * @param y Smoothed y coordinate, in analysis-image pixel space.
 * @param inFrameLikelihood Smoothed detection confidence in [0, 1].
 */
data class SmoothedLandmark(val x: Float, val y: Float, val inFrameLikelihood: Float)

/**
 * @brief Low-pass-filters ML Kit's per-frame [Pose] landmarks with an
 *        exponential moving average, so the drawn skeleton doesn't visibly
 *        jitter/flicker from frame-to-frame detector noise.
 *
 * This smooths `inFrameLikelihood` too, not just position -- without that,
 * a landmark hovering right around [PoseSkeletonRenderer.MIN_LIKELIHOOD]
 * makes whole bones repeatedly pop in and out, which reads as flicker just
 * as much as position jitter does.
 *
 * State is per-landmark-type and carries across calls to [smooth], so this
 * is meant as one instance per detection stream (i.e. per [PoseAnalyzer]) --
 * create a new one whenever the stream restarts rather than reusing one
 * across unrelated streams, or the first frame of the new stream will lerp
 * in from the old stream's last pose.
 *
 * @param smoothingFactor Weight given to each new sample; lower = smoother
 *        but more lag behind the true position. 0.4 noticeably cuts jitter
 *        while still keeping up with a fast bowling arm swing.
 */
class PoseLandmarkSmoother(
    private val smoothingFactor: Float = 0.4f,
) {
    private val previous = mutableMapOf<Int, SmoothedLandmark>()

    /**
     * @brief Applies one frame of exponential smoothing to every landmark
     *        in [pose] and returns the updated smoothed state.
     * @param pose The raw ML Kit detection result for the current frame.
     * @return The smoothed landmarks seen so far, keyed by ML Kit's
     *         `PoseLandmark` type constant (e.g. `PoseLandmark.LEFT_ELBOW`).
     */
    fun smooth(pose: Pose): Map<Int, SmoothedLandmark> {
        for (landmark in pose.allPoseLandmarks) {
            val prev = previous[landmark.landmarkType]
            val next = if (prev == null) {
                SmoothedLandmark(landmark.position.x, landmark.position.y, landmark.inFrameLikelihood)
            } else {
                SmoothedLandmark(
                    x = prev.x + (smoothingFactor * (landmark.position.x - prev.x)),
                    y = prev.y + (smoothingFactor * (landmark.position.y - prev.y)),
                    inFrameLikelihood = landmark.inFrameLikelihood,
                )
            }
            previous[landmark.landmarkType] = next
        }
        return previous.toMap()
    }
}