package com.visionCameraImageLabelerImpreciseImprecise;

import android.annotation.SuppressLint;
import android.media.Image;
import android.util.SparseIntArray;

import androidx.annotation.NonNull;
import androidx.camera.core.ImageProxy;

import com.facebook.react.bridge.WritableNativeArray;
import com.facebook.react.bridge.WritableNativeMap;
import com.google.android.gms.tasks.OnFailureListener;
import com.google.android.gms.tasks.OnSuccessListener;
import com.google.android.gms.tasks.Task;
import com.google.android.gms.tasks.Tasks;
import com.google.mlkit.vision.common.InputImage;
import com.google.mlkit.vision.pose.Pose;
import com.google.mlkit.vision.pose.PoseDetection;
import com.google.mlkit.vision.pose.PoseDetector;
import com.google.mlkit.vision.pose.PoseLandmark;
import com.google.mlkit.vision.pose.defaults.PoseDetectorOptions;
import com.mrousavy.camera.frameprocessor.FrameProcessorPlugin;

import org.jetbrains.annotations.NotNull;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class VisionCameraImageLabelerImprecisePlugin extends FrameProcessorPlugin {
    private final PoseDetectorOptions options =
            new PoseDetectorOptions.Builder()
                    .setDetectorMode(PoseDetectorOptions.STREAM_MODE)
                    .build();

    private final PoseDetector poseDetector = PoseDetection.getClient(options);
    private static final SparseIntArray ORIENTATIONS = new SparseIntArray();


    @Override
    public Object callback(ImageProxy frame, @NotNull Object[] params) {
        @SuppressLint("UnsafeOptInUsageError")
        Image mediaImage = frame.getImage();
        if (mediaImage != null) {
            InputImage image = InputImage.fromMediaImage(mediaImage, frame.getImageInfo().getRotationDegrees());

            Task<Pose> task = poseDetector.process(image);

   
            try {
                Pose pose = Tasks.await(task);

                WritableNativeMap positions = new WritableNativeMap();

                PoseLandmark leftShoulder = pose.getPoseLandmark(PoseLandmark.LEFT_SHOULDER);
                double LEFT_SHOULDER_X = leftShoulder.getPosition3D().getX();
                double LEFT_SHOULDER_Y = leftShoulder.getPosition3D().getY();
                positions.putDouble("LEFT_SHOULDER_X", LEFT_SHOULDER_X);
                positions.putDouble("LEFT_SHOULDER_Y", LEFT_SHOULDER_Y);
                PoseLandmark rightShoulder = pose.getPoseLandmark(PoseLandmark.RIGHT_SHOULDER);
                double RIGHT_SHOULDER_X = rightShoulder.getPosition3D().getX();
                double RIGHT_SHOULDER_Y = rightShoulder.getPosition3D().getY();
                positions.putDouble("RIGHT_SHOULDER_X", RIGHT_SHOULDER_X);
                positions.putDouble("RIGHT_SHOULDER_Y", RIGHT_SHOULDER_Y);
                PoseLandmark leftElbow = pose.getPoseLandmark(PoseLandmark.LEFT_ELBOW);
                double LEFT_ELBOW_X = leftElbow.getPosition3D().getX();
                double LEFT_ELBOW_Y = leftElbow.getPosition3D().getY();
                positions.putDouble("LEFT_ELBOW_X", LEFT_ELBOW_X);
                positions.putDouble("LEFT_ELBOW_Y", LEFT_ELBOW_Y);
                PoseLandmark rightElbow = pose.getPoseLandmark(PoseLandmark.RIGHT_ELBOW);
                double RIGHT_ELBOW_X = rightElbow.getPosition3D().getX();
                double RIGHT_ELBOW_Y = rightElbow.getPosition3D().getY();
                positions.putDouble("RIGHT_ELBOW_X", RIGHT_ELBOW_X);
                positions.putDouble("RIGHT_ELBOW_Y", RIGHT_ELBOW_Y);
                PoseLandmark leftWrist = pose.getPoseLandmark(PoseLandmark.LEFT_WRIST);
                double LEFT_WRIST_X = leftWrist.getPosition3D().getX();
                double LEFT_WRIST_Y = leftWrist.getPosition3D().getY();
                positions.putDouble("LEFT_WRIST_X", LEFT_WRIST_X);
                positions.putDouble("LEFT_WRIST_Y", LEFT_WRIST_Y);
                PoseLandmark rightWrist = pose.getPoseLandmark(PoseLandmark.RIGHT_WRIST);
                double RIGHT_WRIST_X = rightWrist.getPosition3D().getX();
                double RIGHT_WRIST_Y = rightWrist.getPosition3D().getY();
                positions.putDouble("RIGHT_WRIST_X", RIGHT_WRIST_X);
                positions.putDouble("RIGHT_WRIST_Y", RIGHT_WRIST_Y);
                PoseLandmark leftHip = pose.getPoseLandmark(PoseLandmark.LEFT_HIP);
                double LEFT_HIP_X = leftHip.getPosition3D().getX();
                double LEFT_HIP_Y = leftHip.getPosition3D().getY();
                positions.putDouble("LEFT_HIP_X", LEFT_HIP_X);
                positions.putDouble("LEFT_HIP_Y", LEFT_HIP_Y);
                PoseLandmark rightHip = pose.getPoseLandmark(PoseLandmark.RIGHT_HIP);
                double RIGHT_HIP_X = rightHip.getPosition3D().getX();
                double RIGHT_HIP_Y = rightHip.getPosition3D().getY();
                positions.putDouble("RIGHT_HIP_X", RIGHT_HIP_X);
                positions.putDouble("RIGHT_HIP_Y", RIGHT_HIP_Y);
                PoseLandmark leftKnee = pose.getPoseLandmark(PoseLandmark.LEFT_KNEE);
                double LEFT_KNEE_X = leftKnee.getPosition3D().getX();
                double LEFT_KNEE_Y = leftKnee.getPosition3D().getY();
                positions.putDouble("LEFT_KNEE_X", LEFT_KNEE_X);
                positions.putDouble("LEFT_KNEE_Y", LEFT_KNEE_Y);
                PoseLandmark rightKnee = pose.getPoseLandmark(PoseLandmark.RIGHT_KNEE);
                double RIGHT_KNEE_X = rightKnee.getPosition3D().getX();
                double RIGHT_KNEE_Y = rightKnee.getPosition3D().getY();
                positions.putDouble("RIGHT_KNEE_X", RIGHT_KNEE_X);
                positions.putDouble("RIGHT_KNEE_Y", RIGHT_KNEE_Y);
                PoseLandmark leftHeel = pose.getPoseLandmark(PoseLandmark.LEFT_HEEL);
                double LEFT_HEEL_X = leftHeel.getPosition3D().getX();
                double LEFT_HEEL_Y = leftHeel.getPosition3D().getY();
                positions.putDouble("LEFT_HEEL_X", LEFT_HEEL_X);
                positions.putDouble("LEFT_HEEL_Y", LEFT_HEEL_Y);
                PoseLandmark rightHeel = pose.getPoseLandmark(PoseLandmark.RIGHT_HEEL);
                double RIGHT_HEEL_X = rightHeel.getPosition3D().getX();
                double RIGHT_HEEL_Y = rightHeel.getPosition3D().getY();
                positions.putDouble("RIGHT_HEEL_X", RIGHT_HEEL_X);
                positions.putDouble("RIGHT_HEEL_Y", RIGHT_HEEL_Y);
                PoseLandmark leftFootIndex = pose.getPoseLandmark(PoseLandmark.LEFT_FOOT_INDEX);
                double LEFT_FOOT_INDEX_X = leftFootIndex.getPosition3D().getX();
                double LEFT_FOOT_INDEX_Y = leftFootIndex.getPosition3D().getY();
                positions.putDouble("LEFT_FOOT_INDEX_X", LEFT_FOOT_INDEX_X);
                positions.putDouble("LEFT_FOOT_INDEX_Y", LEFT_FOOT_INDEX_Y);
                PoseLandmark rightFootIndex = pose.getPoseLandmark(PoseLandmark.RIGHT_FOOT_INDEX);
                double RIGHT_FOOT_INDEX_X = rightFootIndex.getPosition3D().getX();
                double RIGHT_FOOT_INDEX_Y = rightFootIndex.getPosition3D().getY();
                positions.putDouble("RIGHT_FOOT_INDEX_X", RIGHT_FOOT_INDEX_X);
                positions.putDouble("RIGHT_FOOT_INDEX_Y", RIGHT_FOOT_INDEX_Y);
                PoseLandmark nose = pose.getPoseLandmark(PoseLandmark.NOSE);
                double NOSE_X = nose.getPosition3D().getX();
                double NOSE_Y = nose.getPosition3D().getY();
                positions.putDouble("NOSE_X", NOSE_X);
                positions.putDouble("NOSE_Y", NOSE_Y);
                PoseLandmark leftEye = pose.getPoseLandmark(PoseLandmark.LEFT_EYE);
                double LEFT_EYE_X = leftEye.getPosition3D().getX();
                double LEFT_EYE_Y = leftEye.getPosition3D().getY();
                positions.putDouble("LEFT_EYE_X", LEFT_EYE_X);
                positions.putDouble("LEFT_EYE_Y", LEFT_EYE_Y);
                PoseLandmark rightEye = pose.getPoseLandmark(PoseLandmark.RIGHT_EYE);
                double RIGHT_EYE_X = rightEye.getPosition3D().getX();
                double RIGHT_EYE_Y = rightEye.getPosition3D().getY();
                positions.putDouble("RIGHT_EYE_X", RIGHT_EYE_X);
                positions.putDouble("RIGHT_EYE_Y", RIGHT_EYE_Y);
                PoseLandmark leftMouth = pose.getPoseLandmark(PoseLandmark.LEFT_MOUTH);
                double LEFT_MOUTH_X = leftMouth.getPosition3D().getX();
                double LEFT_MOUTH_Y = leftMouth.getPosition3D().getY();
                positions.putDouble("LEFT_MOUTH_X", LEFT_MOUTH_X);
                positions.putDouble("LEFT_MOUTH_Y", LEFT_MOUTH_Y);
                PoseLandmark rightMouth = pose.getPoseLandmark(PoseLandmark.RIGHT_MOUTH);
                double RIGHT_MOUTH_X = rightMouth.getPosition3D().getX();
                double RIGHT_MOUTH_Y = rightMouth.getPosition3D().getY();
                positions.putDouble("RIGHT_MOUTH_X", RIGHT_MOUTH_X);
                positions.putDouble("RIGHT_MOUTH_Y", RIGHT_MOUTH_Y);

                PoseLandmark leftIndex = pose.getPoseLandmark(PoseLandmark.LEFT_INDEX);
                double LEFT_INDEX_X = leftIndex.getPosition3D().getX();
                double LEFT_INDEX_Y = leftIndex.getPosition3D().getY();
                positions.putDouble("LEFT_INDEX_X", LEFT_INDEX_X);
                positions.putDouble("LEFT_INDEX_Y", LEFT_INDEX_Y);
                PoseLandmark rightIndex = pose.getPoseLandmark(PoseLandmark.RIGHT_INDEX);
                double RIGHT_INDEX_X = rightIndex.getPosition3D().getX();
                double RIGHT_INDEX_Y = rightIndex.getPosition3D().getY();
                positions.putDouble("RIGHT_INDEX_X", RIGHT_INDEX_X);
                positions.putDouble("RIGHT_INDEX_Y", RIGHT_INDEX_Y);

                double leftShoulderInFrameLikelihood = leftShoulder.getInFrameLikelihood();
                double rightShoulderInFrameLikelihood = leftShoulder.getInFrameLikelihood();
                double leftElbowInFrameLikelihood = leftShoulder.getInFrameLikelihood();
                double rightElbowInFrameLikelihood = rightElbow.getInFrameLikelihood();
                double leftWristInFrameLikelihood = leftWrist.getInFrameLikelihood();
                double rightWristInFrameLikelihood = rightWrist.getInFrameLikelihood();
                double leftHipInFrameLikelihood = leftHip.getInFrameLikelihood();
                double rightHipInFrameLikelihood = rightHip.getInFrameLikelihood();
                double leftKneeInFrameLikelihood = leftKnee.getInFrameLikelihood();
                double rightKneeInFrameLikelihood = rightKnee.getInFrameLikelihood();
                double leftIndexFingerInFrameLikelihood = leftIndex.getInFrameLikelihood();
                double rightIndexFingerInFrameLikelihood = rightIndex.getInFrameLikelihood();
                double leftHeelInFrameLikelihood = leftHeel.getInFrameLikelihood();
                double rightHeelInFrameLikelihood = rightHeel.getInFrameLikelihood();
                double leftToeInFrameLikelihood = leftFootIndex.getInFrameLikelihood();
                double rightToeInFrameLikelihood = rightFootIndex.getInFrameLikelihood();
                double noseInFrameLikelihood = nose.getInFrameLikelihood();
                double leftEyeInFrameLikelihood = leftEye.getInFrameLikelihood();
                double rightEyeInFrameLikelihood = rightEye.getInFrameLikelihood();
                double mouthLeftInFrameLikelihood = leftMouth.getInFrameLikelihood();
                double mouthRightInFrameLikelihood = rightMouth.getInFrameLikelihood();

                double meanLikelihood = 0.04761904761 * (
                    leftShoulderInFrameLikelihood +
                    rightShoulderInFrameLikelihood +
                    leftElbowInFrameLikelihood +
                    rightElbowInFrameLikelihood +
                    leftWristInFrameLikelihood +
                    rightWristInFrameLikelihood +
                    leftHipInFrameLikelihood +
                    rightHipInFrameLikelihood +
                    leftKneeInFrameLikelihood +
                    rightKneeInFrameLikelihood +
                    leftIndexFingerInFrameLikelihood +
                    rightIndexFingerInFrameLikelihood +
                    leftHeelInFrameLikelihood +
                    rightHeelInFrameLikelihood +
                    leftToeInFrameLikelihood +
                    rightToeInFrameLikelihood +
                    noseInFrameLikelihood +
                    leftEyeInFrameLikelihood +
                    rightEyeInFrameLikelihood +
                    mouthLeftInFrameLikelihood +
                    mouthRightInFrameLikelihood
                );
                positions.putDouble("MEAN_IN_FRAME_LIKELIHOOD", meanLikelihood);

                return positions;
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return null;
    }

    VisionCameraImageLabelerImprecisePlugin() {
        super("labelImageImprecise");
    }
}