//
//  VisionCameraImageLabelerImprecise.m
//  VisionCameraExample
//
//  Created by Marc Rousavy on 06.05.21.
//

#import <MLKit.h>
#import <VisionCamera/Frame.h>
#import <VisionCamera/FrameProcessorPlugin.h>

// Example for an Objective-C Frame Processor plugin

@interface QRCodeFrameProcessorPluginObjC : NSObject

//+ (MLKImageLabeler *)labeler;
+ (MLKPoseDetector *)detector;

@end

@implementation QRCodeFrameProcessorPluginObjC

+ (MLKPoseDetector *)detector {
  static MLKPoseDetector *detector = nil;
  if (detector == nil) {
    MLKAccuratePoseDetectorOptions *options =
        [[MLKAccuratePoseDetectorOptions alloc] init];
    options.detectorMode = MLKPoseDetectorModeSingleImage;
    detector = [MLKPoseDetector poseDetectorWithOptions:options];
  }
  return detector;
}

static inline id labelImageImprecise(Frame *frame, NSArray *arguments) {
  MLKVisionImage *image = [[MLKVisionImage alloc] initWithBuffer:frame.buffer];
  image.orientation = frame.orientation; // <-- TODO: is mirrored?

  NSError *error;
  NSArray<MLKPose *> *poses =
      [[QRCodeFrameProcessorPluginObjC detector] resultsInImage:image
                                                          error:&error];

  NSMutableArray *results = [NSMutableArray arrayWithCapacity:poses.count];
  for (MLKPose *pose in poses) {

    /** The landmark which corresponds to the nose. */
    MLKPoseLandmark *nose = [pose landmarkOfType:MLKPoseLandmarkTypeNose];
    MLKVision3DPoint *nosePosition = nose.position;

    /** The landmark which corresponds to the left eye inner edge. */
    MLKPoseLandmark *leftEyeInner =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftEyeInner];
    MLKVision3DPoint *leftEyeInnerPosition = leftEyeInner.position;

    /** The landmark which corresponds to the left eye. */
    MLKPoseLandmark *leftEye = [pose landmarkOfType:MLKPoseLandmarkTypeLeftEye];
    MLKVision3DPoint *leftEyePosition = leftEye.position;

    /** The landmark which corresponds to the left eye outer edge. */
    MLKPoseLandmark *leftEyeOuter =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftEyeOuter];
    MLKVision3DPoint *leftEyeOuterPosition = leftEyeOuter.position;

    /** The landmark which corresponds to the right eye inner edge. */
    MLKPoseLandmark *rightEyeInner =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightEyeInner];
    MLKVision3DPoint *rightEyeInnerPosition = rightEyeInner.position;

    /** The landmark which corresponds to the right eye. */
    MLKPoseLandmark *rightEye =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightEye];
    MLKVision3DPoint *rightEyePosition = rightEye.position;

    /** The landmark which corresponds to the right eye outer edge. */
    MLKPoseLandmark *rightEyeOuter =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightEyeOuter];
    MLKVision3DPoint *rightEyeOuterPosition = rightEyeOuter.position;

    /** The landmark which corresponds to the left ear. */
    MLKPoseLandmark *leftEar = [pose landmarkOfType:MLKPoseLandmarkTypeLeftEar];
    MLKVision3DPoint *leftEarPosition = leftEar.position;

    /** The landmark which corresponds to the right ear. */
    MLKPoseLandmark *rightEar =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightEar];
    MLKVision3DPoint *rightEarPosition = rightEar.position;

    /** The landmark which corresponds to the left mouth edge. */
    MLKPoseLandmark *mouthLeft =
        [pose landmarkOfType:MLKPoseLandmarkTypeMouthLeft];
    MLKVision3DPoint *mouthLeftPosition = mouthLeft.position;

    /** The landmark which corresponds to the right mouth edge. */
    MLKPoseLandmark *mouthRight =
        [pose landmarkOfType:MLKPoseLandmarkTypeMouthRight];
    MLKVision3DPoint *mouthRightPosition = mouthRight.position;

    /** The landmark which corresponds to the left shoulder. */
    MLKPoseLandmark *leftShoulder =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftShoulder];
    MLKVision3DPoint *leftShoulderPosition = leftShoulder.position;

    /** The landmark which corresponds to the right shoulder. */
    MLKPoseLandmark *rightShoulder =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightShoulder];
    MLKVision3DPoint *rightShoulderPosition = rightShoulder.position;

    /** The landmark which corresponds to the left elbow. */
    MLKPoseLandmark *leftElbow =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftElbow];
    MLKVision3DPoint *leftElbowPosition = leftElbow.position;

    /** The landmark which corresponds to the right elbow. */
    MLKPoseLandmark *rightElbow =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightElbow];
    MLKVision3DPoint *rightElbowPosition = rightElbow.position;

    /** The landmark which corresponds to the left wrist. */
    MLKPoseLandmark *leftWrist =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftWrist];
    MLKVision3DPoint *leftWristPosition = leftWrist.position;

    /** The landmark which corresponds to the right wrist. */
    MLKPoseLandmark *rightWrist =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightWrist];
    MLKVision3DPoint *rightWristPosition = rightWrist.position;

    /** The landmark which corresponds to the pinky finger on the left hand. */
    MLKPoseLandmark *leftPinkyFinger =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftPinkyFinger];
    MLKVision3DPoint *leftPinkyFingerPosition = leftPinkyFinger.position;

    /** The landmark which corresponds to the pinky finger on the right hand. */
    MLKPoseLandmark *rightPinkyFinger =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightPinkyFinger];
    MLKVision3DPoint *rightPinkyFingerPosition = rightPinkyFinger.position;

    /** The landmark which corresponds to the index finger on the left hand. */
    MLKPoseLandmark *leftIndexFinger =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftIndexFinger];
    MLKVision3DPoint *leftIndexFingerPosition = leftIndexFinger.position;

    /** The landmark which corresponds to the index finger on the right hand. */
    MLKPoseLandmark *rightIndexFinger =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightIndexFinger];
    MLKVision3DPoint *rightIndexFingerPosition = rightIndexFinger.position;

    /** The landmark which corresponds to the left thumb. */
    MLKPoseLandmark *leftThumb =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftThumb];
    MLKVision3DPoint *leftThumbPosition = leftThumb.position;

    /** The landmark which corresponds to the right thumb. */
    MLKPoseLandmark *rightThumb =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightThumb];
    MLKVision3DPoint *rightThumbPosition = rightThumb.position;

    /** The landmark which corresponds to the left hip. */
    MLKPoseLandmark *leftHip = [pose landmarkOfType:MLKPoseLandmarkTypeLeftHip];
    MLKVision3DPoint *leftHipPosition = leftHip.position;

    /** The landmark which corresponds to the right hip. */
    MLKPoseLandmark *rightHip =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightHip];
    MLKVision3DPoint *rightHipPosition = rightHip.position;

    /** The landmark which corresponds to the left knee. */
    MLKPoseLandmark *leftKnee =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftKnee];
    MLKVision3DPoint *leftKneePosition = leftKnee.position;

    /** The landmark which corresponds to the right knee. */
    MLKPoseLandmark *rightKnee =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightKnee];
    MLKVision3DPoint *rightKneePosition = rightKnee.position;

    /** The landmark which corresponds to the left ankle. */
    MLKPoseLandmark *leftAnkle =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftAnkle];
    MLKVision3DPoint *leftAnklePosition = leftAnkle.position;

    /** The landmark which corresponds to the right ankle. */
    MLKPoseLandmark *rightAnkle =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightAnkle];
    MLKVision3DPoint *rightAnklePosition = rightAnkle.position;

    /** The landmark which corresponds to the left heel. */
    MLKPoseLandmark *leftHeel =
        [pose landmarkOfType:MLKPoseLandmarkTypeLeftHeel];
    MLKVision3DPoint *leftHeelPosition = leftHeel.position;

    /** The landmark which corresponds to the right heel. */
    MLKPoseLandmark *rightHeel =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightHeel];
    MLKVision3DPoint *rightHeelPosition = rightHeel.position;

    /** The landmark which corresponds to the toe on the left foot. */
    MLKPoseLandmark *leftToe = [pose landmarkOfType:MLKPoseLandmarkTypeLeftToe];
    MLKVision3DPoint *leftToePosition = leftToe.position;

    /** The landmark which corresponds to the toe on the right foot. */
    MLKPoseLandmark *rightToe =
        [pose landmarkOfType:MLKPoseLandmarkTypeRightToe];
    MLKVision3DPoint *rightToePosition = rightToe.position;

    float meanInFrameLikelihood =
        0.04761904761 *
        (leftEyeInner.inFrameLikelihood + leftEye.inFrameLikelihood +
         rightEye.inFrameLikelihood + mouthLeft.inFrameLikelihood +
         mouthRight.inFrameLikelihood + leftShoulder.inFrameLikelihood +
         rightShoulder.inFrameLikelihood + leftElbow.inFrameLikelihood +
         rightElbow.inFrameLikelihood + leftWrist.inFrameLikelihood +
         rightWrist.inFrameLikelihood + leftIndexFinger.inFrameLikelihood +
         rightIndexFinger.inFrameLikelihood + leftHip.inFrameLikelihood +
         rightHip.inFrameLikelihood + leftKnee.inFrameLikelihood +
         rightKnee.inFrameLikelihood + leftHeel.inFrameLikelihood +
         rightHeel.inFrameLikelihood + leftToe.inFrameLikelihood +
         rightToe.inFrameLikelihood);

    [results addObject : @{
      @"NOSE_X" : @(nosePosition.x),
      @"NOSE_Y" : @(nosePosition.y),

      @"LEFT_EYE_X" : @(leftEyePosition.x),
      @"LEFT_EYE_Y" : @(leftEyePosition.y),

      @"RIGHT_EYE_X" : @(rightEyePosition.x),
      @"RIGHT_EYE_Y" : @(rightEyePosition.y),

      @"LEFT_MOUTH_X" : @(mouthLeftPosition.x),
      @"LEFT_MOUTH_Y" : @(mouthLeftPosition.y),

      @"RIGHT_MOUTH_X" : @(mouthRightPosition.x),
      @"RIGHT_MOUTH_Y" : @(mouthRightPosition.y),

      @"LEFT_SHOULDER_X" : @(leftShoulderPosition.x),
      @"LEFT_SHOULDER_Y" : @(leftShoulderPosition.y),

      @"RIGHT_SHOULDER_X" : @(rightShoulderPosition.x),
      @"RIGHT_SHOULDER_Y" : @(rightShoulderPosition.y),

      @"LEFT_ELBOW_X" : @(leftElbowPosition.x),
      @"LEFT_ELBOW_Y" : @(leftElbowPosition.y),

      @"RIGHT_ELBOW_X" : @(rightElbowPosition.x),
      @"RIGHT_ELBOW_Y" : @(rightElbowPosition.y),

      @"LEFT_WRIST_X" : @(leftWristPosition.x),
      @"LEFT_WRIST_Y" : @(leftWristPosition.y),

      @"RIGHT_WRIST_X" : @(rightWristPosition.x),
      @"RIGHT_WRIST_Y" : @(rightWristPosition.y),

      @"LEFT_INDEX_X" : @(leftIndexFingerPosition.x),
      @"LEFT_INDEX_Y" : @(leftIndexFingerPosition.y),

      @"RIGHT_INDEX_X" : @(rightIndexFingerPosition.x),
      @"RIGHT_INDEX_Y" : @(rightIndexFingerPosition.y),

      @"LEFT_HIP_X" : @(leftHipPosition.x),
      @"LEFT_HIP_Y" : @(leftHipPosition.y),

      @"RIGHT_HIP_X" : @(rightHipPosition.x),
      @"RIGHT_HIP_Y" : @(rightHipPosition.y),

      @"LEFT_KNEE_X" : @(leftKneePosition.x),
      @"LEFT_KNEE_Y" : @(leftKneePosition.y),

      @"RIGHT_KNEE_X" : @(rightKneePosition.x),
      @"RIGHT_KNEE_Y" : @(rightKneePosition.y),

      @"LEFT_HEEL_X" : @(leftHeelPosition.x),
      @"LEFT_HEEL_Y" : @(leftHeelPosition.y),

      @"RIGHT_HEEL_X" : @(rightHeelPosition.x),
      @"RIGHT_HEEL_Y" : @(rightHeelPosition.y),

      @"LEFT_FOOT_INDEX_X" : @(leftToePosition.x),
      @"LEFT_FOOT_INDEX_Y" : @(leftToePosition.y),

      @"RIGHT_FOOT_INDEX_X" : @(rightToePosition.x),
      @"RIGHT_FOOT_INDEX_Y" : @(rightToePosition.y),

      @"MEAN_IN_FRAME_LIKELIHOOD" : @(meanInFrameLikelihood)
    }];
  }

  return results;
}

VISION_EXPORT_FRAME_PROCESSOR(labelImageImprecise)

@end
