// 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 "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/graphs/object_detection_3d/calculators/annotations_to_render_data_calculator.pb.h"
#include "mediapipe/modules/objectron/calculators/annotation_data.pb.h"
#include "mediapipe/util/color.pb.h"
#include "mediapipe/util/render_data.pb.h"

namespace mediapipe {

namespace {

constexpr char kAnnotationTag[] = "ANNOTATIONS";
constexpr char kRenderDataTag[] = "RENDER_DATA";
constexpr char kKeypointLabel[] = "KEYPOINT";
constexpr int kMaxLandmarkThickness = 18;

inline void SetColor(RenderAnnotation* annotation, const Color& color) {
  annotation->mutable_color()->set_r(color.r());
  annotation->mutable_color()->set_g(color.g());
  annotation->mutable_color()->set_b(color.b());
}

// Remap x from range [lo hi] to range [0 1] then multiply by scale.
inline float Remap(float x, float lo, float hi, float scale) {
  return (x - lo) / (hi - lo + 1e-6) * scale;
}

inline void GetMinMaxZ(const FrameAnnotation& annotations, float* z_min,
                       float* z_max) {
  *z_min = std::numeric_limits<float>::max();
  *z_max = std::numeric_limits<float>::min();
  // Use a global depth scale for all the objects in the scene
  for (const auto& object : annotations.annotations()) {
    for (const auto& keypoint : object.keypoints()) {
      *z_min = std::min(keypoint.point_2d().depth(), *z_min);
      *z_max = std::max(keypoint.point_2d().depth(), *z_max);
    }
  }
}

void SetColorSizeValueFromZ(float z, float z_min, float z_max,
                            RenderAnnotation* render_annotation) {
  const int color_value = 255 - static_cast<int>(Remap(z, z_min, z_max, 255));
  ::mediapipe::Color color;
  color.set_r(color_value);
  color.set_g(color_value);
  color.set_b(color_value);
  SetColor(render_annotation, color);
  const int thickness = static_cast<int>((1.f - Remap(z, z_min, z_max, 1)) *
                                         kMaxLandmarkThickness);
  render_annotation->set_thickness(thickness);
}

}  // namespace

// A calculator that converts FrameAnnotation proto to RenderData proto for
// visualization. The input should be the FrameAnnotation proto buffer. It is
// also possible to specify the connections between landmarks.
//
// Example config:
// node {
//   calculator: "AnnotationsToRenderDataCalculator"
//   input_stream: "ANNOTATIONS:annotations"
//   output_stream: "RENDER_DATA:render_data"
//   options {
//     [AnnotationsToRenderDataCalculator.ext] {
//       landmark_connections: [0, 1, 1, 2]
//       landmark_color { r: 0 g: 255 b: 0 }
//       connection_color { r: 0 g: 255 b: 0 }
//       thickness: 4.0
//     }
//   }
// }
class AnnotationsToRenderDataCalculator : public CalculatorBase {
 public:
  AnnotationsToRenderDataCalculator() {}
  ~AnnotationsToRenderDataCalculator() override {}
  AnnotationsToRenderDataCalculator(const AnnotationsToRenderDataCalculator&) =
      delete;
  AnnotationsToRenderDataCalculator& operator=(
      const AnnotationsToRenderDataCalculator&) = delete;

  static absl::Status GetContract(CalculatorContract* cc);

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

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

 private:
  static void SetRenderAnnotationColorThickness(
      const AnnotationsToRenderDataCalculatorOptions& options,
      RenderAnnotation* render_annotation);
  static RenderAnnotation* AddPointRenderData(
      const AnnotationsToRenderDataCalculatorOptions& options,
      RenderData* render_data);

  // Add a command to draw a line in the rendering queue. The line is drawn from
  // (start_x, start_y) to (end_x, end_y). The input x,y can either be in pixel
  // or normalized coordinate [0, 1] as indicated by the normalized flag.
  static void AddConnectionToRenderData(
      float start_x, float start_y, float end_x, float end_y,
      const AnnotationsToRenderDataCalculatorOptions& options, bool normalized,
      RenderData* render_data);

  // Same as above function. Instead of using color data to render the line, it
  // re-colors the line according to the two depth value. gray_val1 is the color
  // of the starting point and gray_val2 is the color of the ending point. The
  // line is colored using gradient color from gray_val1 to gray_val2. The
  // gray_val ranges from [0 to 255] for black to white.
  static void AddConnectionToRenderData(
      float start_x, float start_y, float end_x, float end_y,
      const AnnotationsToRenderDataCalculatorOptions& options, bool normalized,
      int gray_val1, int gray_val2, RenderData* render_data);

  AnnotationsToRenderDataCalculatorOptions options_;
};
REGISTER_CALCULATOR(AnnotationsToRenderDataCalculator);

absl::Status AnnotationsToRenderDataCalculator::GetContract(
    CalculatorContract* cc) {
  RET_CHECK(cc->Inputs().HasTag(kAnnotationTag)) << "No input stream found.";
  if (cc->Inputs().HasTag(kAnnotationTag)) {
    cc->Inputs().Tag(kAnnotationTag).Set<FrameAnnotation>();
  }
  cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();

  return absl::OkStatus();
}

absl::Status AnnotationsToRenderDataCalculator::Open(CalculatorContext* cc) {
  cc->SetOffset(TimestampDiff(0));
  options_ = cc->Options<AnnotationsToRenderDataCalculatorOptions>();

  return absl::OkStatus();
}

absl::Status AnnotationsToRenderDataCalculator::Process(CalculatorContext* cc) {
  auto render_data = absl::make_unique<RenderData>();
  bool visualize_depth = options_.visualize_landmark_depth();
  float z_min = 0.f;
  float z_max = 0.f;

  if (cc->Inputs().HasTag(kAnnotationTag)) {
    const auto& annotations =
        cc->Inputs().Tag(kAnnotationTag).Get<FrameAnnotation>();
    RET_CHECK_EQ(options_.landmark_connections_size() % 2, 0)
        << "Number of entries in landmark connections must be a multiple of 2";

    if (visualize_depth) {
      GetMinMaxZ(annotations, &z_min, &z_max);
      // Only change rendering if there are actually z values other than 0.
      visualize_depth &= ((z_max - z_min) > 1e-3);
    }

    for (const auto& object : annotations.annotations()) {
      for (const auto& keypoint : object.keypoints()) {
        auto* keypoint_data_render =
            AddPointRenderData(options_, render_data.get());
        auto* point = keypoint_data_render->mutable_point();
        if (visualize_depth) {
          SetColorSizeValueFromZ(keypoint.point_2d().depth(), z_min, z_max,
                                 keypoint_data_render);
        }

        point->set_normalized(true);
        point->set_x(keypoint.point_2d().x());
        point->set_y(keypoint.point_2d().y());
      }

      // Add edges
      for (int i = 0; i < options_.landmark_connections_size(); i += 2) {
        const auto& ld0 =
            object.keypoints(options_.landmark_connections(i)).point_2d();
        const auto& ld1 =
            object.keypoints(options_.landmark_connections(i + 1)).point_2d();
        const bool normalized = true;

        if (visualize_depth) {
          const int gray_val1 =
              255 - static_cast<int>(Remap(ld0.depth(), z_min, z_max, 255));
          const int gray_val2 =
              255 - static_cast<int>(Remap(ld1.depth(), z_min, z_max, 255));
          AddConnectionToRenderData(ld0.x(), ld0.y(), ld1.x(), ld1.y(),
                                    options_, normalized, gray_val1, gray_val2,
                                    render_data.get());
        } else {
          AddConnectionToRenderData(ld0.x(), ld0.y(), ld1.x(), ld1.y(),
                                    options_, normalized, render_data.get());
        }
      }
    }
  }

  cc->Outputs()
      .Tag(kRenderDataTag)
      .Add(render_data.release(), cc->InputTimestamp());

  return absl::OkStatus();
}

void AnnotationsToRenderDataCalculator::AddConnectionToRenderData(
    float start_x, float start_y, float end_x, float end_y,
    const AnnotationsToRenderDataCalculatorOptions& options, bool normalized,
    int gray_val1, int gray_val2, RenderData* render_data) {
  auto* connection_annotation = render_data->add_render_annotations();
  RenderAnnotation::GradientLine* line =
      connection_annotation->mutable_gradient_line();
  line->set_x_start(start_x);
  line->set_y_start(start_y);
  line->set_x_end(end_x);
  line->set_y_end(end_y);
  line->set_normalized(normalized);
  line->mutable_color1()->set_r(gray_val1);
  line->mutable_color1()->set_g(gray_val1);
  line->mutable_color1()->set_b(gray_val1);
  line->mutable_color2()->set_r(gray_val2);
  line->mutable_color2()->set_g(gray_val2);
  line->mutable_color2()->set_b(gray_val2);
  connection_annotation->set_thickness(options.thickness());
}

void AnnotationsToRenderDataCalculator::AddConnectionToRenderData(
    float start_x, float start_y, float end_x, float end_y,
    const AnnotationsToRenderDataCalculatorOptions& options, bool normalized,
    RenderData* render_data) {
  auto* connection_annotation = render_data->add_render_annotations();
  RenderAnnotation::Line* line = connection_annotation->mutable_line();
  line->set_x_start(start_x);
  line->set_y_start(start_y);
  line->set_x_end(end_x);
  line->set_y_end(end_y);
  line->set_normalized(normalized);
  SetColor(connection_annotation, options.connection_color());
  connection_annotation->set_thickness(options.thickness());
}

RenderAnnotation* AnnotationsToRenderDataCalculator::AddPointRenderData(
    const AnnotationsToRenderDataCalculatorOptions& options,
    RenderData* render_data) {
  auto* landmark_data_annotation = render_data->add_render_annotations();
  landmark_data_annotation->set_scene_tag(kKeypointLabel);
  SetRenderAnnotationColorThickness(options, landmark_data_annotation);
  return landmark_data_annotation;
}

void AnnotationsToRenderDataCalculator::SetRenderAnnotationColorThickness(
    const AnnotationsToRenderDataCalculatorOptions& options,
    RenderAnnotation* render_annotation) {
  SetColor(render_annotation, options.landmark_color());
  render_annotation->set_thickness(options.thickness());
}

}  // namespace mediapipe
