// Copyright 2020 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <cmath>
#include <memory>

#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_file_properties.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/graphs/iris_tracking/calculators/iris_to_depth_calculator.pb.h"

namespace mediapipe {

namespace {

constexpr char kIrisTag[] = "IRIS";
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
constexpr char kFocalLengthPixelTag[] = "FOCAL_LENGTH";
constexpr char kImageFilePropertiesTag[] = "IMAGE_FILE_PROPERTIES";
constexpr char kLeftIrisDepthTag[] = "LEFT_IRIS_DEPTH_MM";
constexpr char kRightIrisDepthTag[] = "RIGHT_IRIS_DEPTH_MM";
constexpr int kNumIrisLandmarksPerEye = 5;
constexpr float kDepthWeightUpdate = 0.1;
// Avergae fixed iris size across human beings.
constexpr float kIrisSizeInMM = 11.8;

inline float GetDepth(float x0, float y0, float x1, float y1) {
  return std::sqrt((x0 - x1) * (x0 - x1) + (y0 - y1) * (y0 - y1));
}

inline float GetLandmarkDepth(const NormalizedLandmark& ld0,
                              const NormalizedLandmark& ld1,
                              const std::pair<int, int>& image_size) {
  return GetDepth(ld0.x() * image_size.first, ld0.y() * image_size.second,
                  ld1.x() * image_size.first, ld1.y() * image_size.second);
}

float CalculateIrisDiameter(const NormalizedLandmarkList& landmarks,
                            const std::pair<int, int>& image_size) {
  const float dist_vert = GetLandmarkDepth(landmarks.landmark(1),
                                           landmarks.landmark(2), image_size);
  const float dist_hori = GetLandmarkDepth(landmarks.landmark(3),
                                           landmarks.landmark(4), image_size);
  return (dist_hori + dist_vert) / 2.0f;
}

float CalculateDepth(const NormalizedLandmark& center, float focal_length,
                     float iris_size, float img_w, float img_h) {
  std::pair<float, float> origin{img_w / 2.f, img_h / 2.f};
  const auto y = GetDepth(origin.first, origin.second, center.x() * img_w,
                          center.y() * img_h);
  const auto x = std::sqrt(focal_length * focal_length + y * y);
  const auto depth = kIrisSizeInMM * x / iris_size;
  return depth;
}

}  // namespace

// Estimates depth from iris to camera given focal length and image size.
//
// Usage example:
// node {
//   calculator: "IrisToDepthCalculator"
//   # A NormalizedLandmarkList contains landmarks for both iris.
//   input_stream: "IRIS:iris_landmarks"
//   input_stream: "IMAGE_SIZE:image_size"
//   # Note: Only one of FOCAL_LENGTH or IMAGE_FILE_PROPERTIES is necessary
//   # to get focal length in pixels. Sending focal length in pixels to
//   # this calculator is optional.
//   input_side_packet: "FOCAL_LENGTH:focal_length_pixel"
//   # OR
//   input_side_packet: "IMAGE_FILE_PROPERTIES:image_file_properties"
//   output_stream: "LEFT_IRIS_DEPTH_MM:left_iris_depth_mm"
//   output_stream: "RIGHT_IRIS_DEPTH_MM:right_iris_depth_mm"
// }
class IrisToDepthCalculator : public CalculatorBase {
 public:
  static absl::Status GetContract(CalculatorContract* cc) {
    cc->Inputs().Tag(kIrisTag).Set<NormalizedLandmarkList>();
    cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();

    // Only one of kFocalLengthPixelTag or kImageFilePropertiesTag must exist
    // if they are present.
    RET_CHECK(!(cc->InputSidePackets().HasTag(kFocalLengthPixelTag) &&
                cc->InputSidePackets().HasTag(kImageFilePropertiesTag)));
    if (cc->InputSidePackets().HasTag(kFocalLengthPixelTag)) {
      cc->InputSidePackets().Tag(kFocalLengthPixelTag).SetAny();
    }
    if (cc->InputSidePackets().HasTag(kImageFilePropertiesTag)) {
      cc->InputSidePackets()
          .Tag(kImageFilePropertiesTag)
          .Set<ImageFileProperties>();
    }
    if (cc->Outputs().HasTag(kLeftIrisDepthTag)) {
      cc->Outputs().Tag(kLeftIrisDepthTag).Set<float>();
    }
    if (cc->Outputs().HasTag(kRightIrisDepthTag)) {
      cc->Outputs().Tag(kRightIrisDepthTag).Set<float>();
    }
    return absl::OkStatus();
  }

  absl::Status Open(CalculatorContext* cc) override;

  absl::Status Process(CalculatorContext* cc) override;

 private:
  float focal_length_pixels_ = -1.f;
  // TODO: Consolidate the logic when switching to input stream for
  // focal length.
  bool compute_depth_from_iris_ = false;
  float smoothed_left_depth_mm_ = -1.f;
  float smoothed_right_depth_mm_ = -1.f;

  void GetLeftIris(const NormalizedLandmarkList& lds,
                   NormalizedLandmarkList* iris);
  void GetRightIris(const NormalizedLandmarkList& lds,
                    NormalizedLandmarkList* iris);
  ::mediapipe::IrisToDepthCalculatorOptions options_;
};
REGISTER_CALCULATOR(IrisToDepthCalculator);

absl::Status IrisToDepthCalculator::Open(CalculatorContext* cc) {
  cc->SetOffset(TimestampDiff(0));
  if (cc->InputSidePackets().HasTag(kFocalLengthPixelTag)) {
#if defined(__APPLE__)
    focal_length_pixels_ = *cc->InputSidePackets()
                                .Tag(kFocalLengthPixelTag)
                                .Get<std::unique_ptr<float>>();
#else
    focal_length_pixels_ =
        cc->InputSidePackets().Tag(kFocalLengthPixelTag).Get<float>();
#endif
    compute_depth_from_iris_ = true;
  } else if (cc->InputSidePackets().HasTag(kImageFilePropertiesTag)) {
    const auto& properties = cc->InputSidePackets()
                                 .Tag(kImageFilePropertiesTag)
                                 .Get<ImageFileProperties>();
    focal_length_pixels_ = properties.focal_length_pixels();
    compute_depth_from_iris_ = true;
  }

  options_ = cc->Options<::mediapipe::IrisToDepthCalculatorOptions>();
  return absl::OkStatus();
}

absl::Status IrisToDepthCalculator::Process(CalculatorContext* cc) {
  // Only process if there's input landmarks.
  if (cc->Inputs().Tag(kIrisTag).IsEmpty()) {
    return absl::OkStatus();
  }

  const auto& iris_landmarks =
      cc->Inputs().Tag(kIrisTag).Get<NormalizedLandmarkList>();
  RET_CHECK_EQ(iris_landmarks.landmark_size(), kNumIrisLandmarksPerEye * 2)
      << "Wrong number of iris landmarks";

  std::pair<int, int> image_size;
  RET_CHECK(!cc->Inputs().Tag(kImageSizeTag).IsEmpty());
  image_size = cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();

  auto left_iris = absl::make_unique<NormalizedLandmarkList>();
  auto right_iris = absl::make_unique<NormalizedLandmarkList>();
  GetLeftIris(iris_landmarks, left_iris.get());
  GetRightIris(iris_landmarks, right_iris.get());

  const auto left_iris_size = CalculateIrisDiameter(*left_iris, image_size);
  const auto right_iris_size = CalculateIrisDiameter(*right_iris, image_size);

#if defined(__APPLE__)
  if (cc->InputSidePackets().HasTag(kFocalLengthPixelTag)) {
    focal_length_pixels_ = *cc->InputSidePackets()
                                .Tag(kFocalLengthPixelTag)
                                .Get<std::unique_ptr<float>>();
  }
#endif

  if (compute_depth_from_iris_ && focal_length_pixels_ > 0) {
    const auto left_depth =
        CalculateDepth(left_iris->landmark(0), focal_length_pixels_,
                       left_iris_size, image_size.first, image_size.second);
    const auto right_depth =
        CalculateDepth(right_iris->landmark(0), focal_length_pixels_,
                       right_iris_size, image_size.first, image_size.second);
    smoothed_left_depth_mm_ =
        smoothed_left_depth_mm_ < 0 || std::isinf(smoothed_left_depth_mm_)
            ? left_depth
            : smoothed_left_depth_mm_ * (1 - kDepthWeightUpdate) +
                  left_depth * kDepthWeightUpdate;
    smoothed_right_depth_mm_ =
        smoothed_right_depth_mm_ < 0 || std::isinf(smoothed_right_depth_mm_)
            ? right_depth
            : smoothed_right_depth_mm_ * (1 - kDepthWeightUpdate) +
                  right_depth * kDepthWeightUpdate;

    if (cc->Outputs().HasTag(kLeftIrisDepthTag)) {
      cc->Outputs()
          .Tag(kLeftIrisDepthTag)
          .AddPacket(MakePacket<float>(smoothed_left_depth_mm_)
                         .At(cc->InputTimestamp()));
    }
    if (cc->Outputs().HasTag(kRightIrisDepthTag)) {
      cc->Outputs()
          .Tag(kRightIrisDepthTag)
          .AddPacket(MakePacket<float>(smoothed_right_depth_mm_)
                         .At(cc->InputTimestamp()));
    }
  }
  return absl::OkStatus();
}

void IrisToDepthCalculator::GetLeftIris(const NormalizedLandmarkList& lds,
                                        NormalizedLandmarkList* iris) {
  // Center, top, bottom, left, right
  *iris->add_landmark() = lds.landmark(options_.left_iris_center_index());
  *iris->add_landmark() = lds.landmark(options_.left_iris_top_index());
  *iris->add_landmark() = lds.landmark(options_.left_iris_bottom_index());
  *iris->add_landmark() = lds.landmark(options_.left_iris_left_index());
  *iris->add_landmark() = lds.landmark(options_.left_iris_right_index());
}

void IrisToDepthCalculator::GetRightIris(const NormalizedLandmarkList& lds,
                                         NormalizedLandmarkList* iris) {
  // Center, top, bottom, left, right
  *iris->add_landmark() = lds.landmark(options_.right_iris_center_index());
  *iris->add_landmark() = lds.landmark(options_.right_iris_top_index());
  *iris->add_landmark() = lds.landmark(options_.right_iris_bottom_index());
  *iris->add_landmark() = lds.landmark(options_.right_iris_left_index());
  *iris->add_landmark() = lds.landmark(options_.right_iris_right_index());
}
}  // namespace mediapipe
