PinPoint/app/src/main/java/com/example/jnicpp/bowling/PoseAnalyzer.kt

/**
 * @file PoseAnalyzer.kt
 * @brief CameraX ImageAnalysis.Analyzer that bridges frames into ML Kit's streaming pose detector.
 */
package com.example.jnicpp.bowling

import androidx.annotation.OptIn
import androidx.camera.core.ExperimentalGetImage
import androidx.camera.core.ImageAnalysis
import androidx.camera.core.ImageProxy
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.pose.PoseDetection
import com.google.mlkit.vision.pose.PoseDetector
import com.google.mlkit.vision.pose.accurate.AccuratePoseDetectorOptions

/**
 * @brief Bridges CameraX's [ImageAnalysis] frame stream into ML Kit's
 *        streaming pose detector.
 *
 * CameraX invokes [analyze] on whatever executor was passed to
 * `ImageAnalysis.setAnalyzer(executor, this)` -- as long as that's a
 * background executor (see [CameraXController]), the actual inference work
 * never runs on the UI thread, so it can't block the UI or the video
 * recording pipeline. Note that [onResult]/[onError] themselves fire back on
 * the *main* thread: ML Kit's `Task#addOnSuccessListener`/`addOnFailureListener`
 * without an explicit `Executor` deliver on the main application thread by
 * default, regardless of which thread called `.process()`. That's
 * intentional here -- it means callers (see [BowlingCameraActivity]) can
 * update views directly from [onResult] with no extra thread hop.
 * `STREAM_MODE` on the detector itself also makes ML Kit assume frames
 * arrive close together and reuse state between them, which is what makes
 * it track a moving body smoothly instead of re-detecting from scratch.
 *
 * @param isFrontCamera Callback returning whether the currently-bound
 *        camera is front-facing, consulted fresh for each frame.
 * @param onResult Invoked on the main thread with each frame's detection result.
 * @param onError Invoked on the main thread if ML Kit fails to process a frame.
 */
class PoseAnalyzer(
    private val isFrontCamera: () -> Boolean,
    private val onResult: (PoseFrameResult) -> Unit,
    private val onError: (Exception) -> Unit
) : ImageAnalysis.Analyzer {

    /**
     * @brief Everything [PoseOverlayView] needs to both draw a pose and
     *        correctly map it from analysis-image pixels to view pixels,
     *        plus the joint angles derived from that same pose.
     *
     * See [PoseAngleCalculator] for the angles, so downstream consumers
     * (overlay text, frame-by-frame logging) don't need to recompute them
     * from [landmarks] themselves. [landmarks] is already smoothed across
     * frames (see [PoseLandmarkSmoother]) rather than raw ML Kit output, so
     * anything drawn straight from it doesn't need to smooth it again.
     *
     * @param landmarks Smoothed landmarks for this frame, keyed by ML Kit's
     *        `PoseLandmark` type constant.
     * @param imageWidth Raw analysis-image buffer width, pre-rotation.
     * @param imageHeight Raw analysis-image buffer height, pre-rotation.
     * @param rotationDegrees Rotation needed to bring the buffer upright.
     * @param isFrontCamera Whether this frame came from the front camera.
     * @param angles Joint angles derived from [landmarks] for this frame.
     */
    data class PoseFrameResult(
        val landmarks: Map<Int, SmoothedLandmark>,
        val imageWidth: Int,
        val imageHeight: Int,
        val rotationDegrees: Int,
        val isFrontCamera: Boolean,
        val angles: PoseAngles
    )

    private val detector: PoseDetector = PoseDetection.getClient(
        AccuratePoseDetectorOptions.Builder()
            .setDetectorMode(AccuratePoseDetectorOptions.STREAM_MODE)
            .build()
    )

    // One smoother per analyzer instance -- a fresh PoseAnalyzer (see
    // CameraXController.bindToLifecycle) means a fresh detection stream, so
    // its smoothing state should start clean rather than lerping in from
    // whatever pose the previous stream last saw.
    private val landmarkSmoother = PoseLandmarkSmoother()

    // STRATEGY_KEEP_ONLY_LATEST on the ImageAnalysis use case (see
    // CameraXController) already ensures we're never handed a backlog, but
    // this guards against overlapping calls if the detector ever falls
    // behind the frame producer.
    @Volatile
    private var isProcessing = false

    /**
     * @brief CameraX entry point: called once per analyzed frame with the
     *        latest camera buffer.
     *
     * Drops the frame immediately (after closing it) if a detection call is
     * already in flight or the buffer has no backing image. Otherwise hands
     * the frame to ML Kit asynchronously; [onResult]/[onError] fire later
     * from the detector's own callback once inference completes.
     *
     * @param imageProxy The camera frame to analyze. Always closed before
     *        this function returns or before the async detector call
     *        completes, whichever applies -- CameraX stalls the analysis
     *        pipeline otherwise.
     */
    @OptIn(ExperimentalGetImage::class)
    override fun analyze(imageProxy: ImageProxy) {
        val mediaImage = imageProxy.image
        if (mediaImage == null || isProcessing) {
            imageProxy.close()
            return
        }
        isProcessing = true

        // Capture these before handing off to the async detector call --
        // imageProxy itself is closed as soon as the detector is done with
        // the underlying buffer, so nothing here should read from it after
        // that point.
        val rotationDegrees = imageProxy.imageInfo.rotationDegrees
        val width = imageProxy.width
        val height = imageProxy.height
        val frontCamera = isFrontCamera()

        val inputImage = InputImage.fromMediaImage(mediaImage, rotationDegrees)

        detector.process(inputImage)
            .addOnSuccessListener { pose ->
                val landmarks = landmarkSmoother.smooth(pose)
                onResult(
                    PoseFrameResult(
                        landmarks = landmarks,
                        imageWidth = width,
                        imageHeight = height,
                        rotationDegrees = rotationDegrees,
                        isFrontCamera = frontCamera,
                        // Cheap (four atan2 pairs at most), safe to compute
                        // on every frame right alongside the detection result.
                        angles = PoseAngleCalculator.compute(landmarks)
                    )
                )
            }
            .addOnFailureListener { e -> onError(e) }
            .addOnCompleteListener {
                isProcessing = false
                // Must always close, or CameraX stalls the analysis
                // pipeline waiting for this frame's buffer to be released.
                imageProxy.close()
            }
    }

    /** @brief Releases the underlying ML Kit detector. Call when this analyzer's stream is done. */
    fun close() {
        detector.close()
    }
}