// Copyright 2019 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 "mediapipe/framework/formats/location.h"

#include <algorithm>
#include <cmath>
#include <memory>

#include "absl/memory/memory.h"
#include "absl/strings/substitute.h"
#include "mediapipe/framework/formats/annotation/locus.pb.h"
#include "mediapipe/framework/formats/annotation/rasterization.pb.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/point2.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/framework/tool/status_util.h"
#include "mediapipe/framework/type_map.h"

namespace mediapipe {

namespace {
// Extracts from the BinaryMask, stored as Rasterization in
// the location_data, the tightest bounding box, that contains all pixels
// encoded in the rasterizations.
Rectangle_i MaskToRectangle(const LocationData& location_data) {
  CHECK(location_data.mask().has_rasterization());
  const auto& rasterization = location_data.mask().rasterization();
  if (rasterization.interval_size() == 0) {
    return Rectangle_i(0, 0, 0, 0);
  }
  int xmin = std::numeric_limits<int>::max();
  int xmax = std::numeric_limits<int>::min();
  int ymin = std::numeric_limits<int>::max();
  int ymax = std::numeric_limits<int>::min();
  for (const auto& interval : rasterization.interval()) {
    xmin = std::min(xmin, interval.left_x());
    xmax = std::max(xmax, interval.right_x());
    ymin = std::min(ymin, interval.y());
    ymax = std::max(ymax, interval.y());
  }
  return Rectangle_i(xmin, ymin, xmax - xmin + 1, ymax - ymin + 1);
}

}  // namespace

Location::Location() {}

Location::Location(const LocationData& location_data)
    : location_data_(location_data) {
  CHECK(IsValidLocationData(location_data_));
}

Location Location::CreateGlobalLocation() {
  LocationData location_data;
  location_data.set_format(LocationData::GLOBAL);
  return Location(location_data);
}

Location Location::CreateBBoxLocation(int xmin, int ymin, int width,
                                      int height) {
  LocationData location_data;
  location_data.set_format(LocationData::BOUNDING_BOX);
  auto* bounding_box = location_data.mutable_bounding_box();
  bounding_box->set_xmin(xmin);
  bounding_box->set_ymin(ymin);
  bounding_box->set_width(width);
  bounding_box->set_height(height);
  return Location(location_data);
}

Location Location::CreateBBoxLocation(const Rectangle_i& rect) {
  return CreateBBoxLocation(rect.xmin(), rect.ymin(), rect.Width(),
                            rect.Height());
}

Location Location::CreateBBoxLocation(const ::mediapipe::BoundingBox& bbox) {
  return CreateBBoxLocation(bbox.left_x(), bbox.upper_y(),
                            bbox.right_x() - bbox.left_x(),
                            bbox.lower_y() - bbox.upper_y());
}

Location Location::CreateRelativeBBoxLocation(float relative_xmin,
                                              float relative_ymin,
                                              float relative_width,
                                              float relative_height) {
  LocationData location_data;
  location_data.set_format(LocationData::RELATIVE_BOUNDING_BOX);
  auto* bounding_box = location_data.mutable_relative_bounding_box();
  bounding_box->set_xmin(relative_xmin);
  bounding_box->set_ymin(relative_ymin);
  bounding_box->set_width(relative_width);
  bounding_box->set_height(relative_height);
  return Location(location_data);
}

Location Location::CreateRelativeBBoxLocation(const Rectangle_f& rect) {
  return CreateRelativeBBoxLocation(rect.xmin(), rect.ymin(), rect.Width(),
                                    rect.Height());
}

LocationData::Format Location::GetFormat() const {
  return location_data_.format();
}

// static
bool Location::IsValidLocationData(const LocationData& location_data) {
  switch (location_data.format()) {
    case LocationData::GLOBAL: {
      // Nothing to check for global location data.
      return true;
    }
    case LocationData::BOUNDING_BOX: {
      return location_data.has_bounding_box() &&
             location_data.bounding_box().has_xmin() &&
             location_data.bounding_box().has_ymin() &&
             location_data.bounding_box().has_width() &&
             location_data.bounding_box().has_height();
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      return location_data.has_relative_bounding_box() &&
             location_data.relative_bounding_box().has_xmin() &&
             location_data.relative_bounding_box().has_ymin() &&
             location_data.relative_bounding_box().has_width() &&
             location_data.relative_bounding_box().has_height();
    }
    case LocationData::MASK: {
      return location_data.has_mask() && location_data.mask().has_width() &&
             location_data.mask().has_height() &&
             location_data.mask().has_rasterization();
    }
    default: {
      return false;
    }
  }
}

template <>
Rectangle_i Location::GetBBox<Rectangle_i>() const {
  CHECK_EQ(LocationData::BOUNDING_BOX, location_data_.format());
  const auto& box = location_data_.bounding_box();
  return Rectangle_i(box.xmin(), box.ymin(), box.width(), box.height());
}

Location& Location::Scale(const float scale) {
  CHECK(!location_data_.has_mask())
      << "Location mask scaling is not implemented.";
  CHECK_GT(scale, 0.0f);
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      // Do nothing.
      break;
    }
    case LocationData::BOUNDING_BOX: {
      auto* bb = location_data_.mutable_bounding_box();
      bb->set_xmin(scale * bb->xmin());
      bb->set_ymin(scale * bb->ymin());
      bb->set_width(scale * bb->width());
      bb->set_height(scale * bb->height());
      break;
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      auto* bb = location_data_.mutable_relative_bounding_box();
      bb->set_xmin(scale * bb->xmin());
      bb->set_ymin(scale * bb->ymin());
      bb->set_width(scale * bb->width());
      bb->set_height(scale * bb->height());
      for (auto& keypoint : *location_data_.mutable_relative_keypoints()) {
        keypoint.set_x(scale * keypoint.x());
        keypoint.set_y(scale * keypoint.y());
      }
      break;
    }
    case LocationData::MASK: {
      LOG(FATAL) << "Scaling for location data of type MASK is not supported.";
      break;
    }
  }
  return *this;
}

Location& Location::Square(int image_width, int image_height) {
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      // Do nothing.
      break;
    }
    case LocationData::BOUNDING_BOX: {
      auto* box = location_data_.mutable_bounding_box();
      const int max_dimension = std::max(box->width(), box->height());
      if (max_dimension > box->width()) {
        box->set_xmin(box->xmin() + box->width() / 2 - max_dimension / 2);
        box->set_width(max_dimension);
      } else if (max_dimension > box->height()) {
        box->set_ymin(box->ymin() + box->height() / 2 - max_dimension / 2);
        box->set_height(max_dimension);
      }
      break;
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      auto* box = location_data_.mutable_relative_bounding_box();
      const float absolute_xmin = box->xmin() * image_width;
      const float absolute_ymin = box->ymin() * image_height;
      const float absolute_width = box->width() * image_width;
      const float absolute_height = box->height() * image_height;
      const float max_dimension = std::max(absolute_width, absolute_height);
      if (max_dimension > absolute_width) {
        box->set_xmin((absolute_xmin + absolute_width / 2 - max_dimension / 2) /
                      image_width);
        box->set_width(max_dimension / image_width);
      } else if (max_dimension > absolute_height) {
        box->set_ymin(
            (absolute_ymin + absolute_height / 2 - max_dimension / 2) /
            image_height);
        box->set_height(max_dimension / image_height);
      }
      break;
    }
    case LocationData::MASK: {
      LOG(FATAL) << "Squaring for location data of type MASK is not supported.";
      break;
    }
  }
  return *this;
}

namespace {

// Finds an optimal shift t such that I = [min_value + t, max_value + t) will be
// included in the interval J = [0, range) if possible. If the above is not
// possible, then interval I will be centered at the center of interval J.
// This function is inteded to shift boundaries of intervals such that they
// best fit within an image.
float BestShift(float min_value, float max_value, float range) {
  CHECK_LE(min_value, max_value);
  const float value_range = max_value - min_value;
  if (value_range > range) {
    return 0.5f * (range - min_value - max_value);
  } else {
    if (min_value < 0.0f) {
      return -min_value;
    } else if (max_value > range) {
      return range - max_value;
    }
  }
  return 0.0f;
}

}  // namespace

Location& Location::ShiftToFitBestIntoImage(int image_width, int image_height) {
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      // Do nothing.
      break;
    }
    case LocationData::BOUNDING_BOX: {
      auto* box = location_data_.mutable_bounding_box();
      box->set_xmin(static_cast<int>(std::round(
          box->xmin() +
          BestShift(box->xmin(), box->xmin() + box->width(), image_width))));
      box->set_ymin(static_cast<int>(std::round(
          box->ymin() +
          BestShift(box->ymin(), box->ymin() + box->height(), image_height))));
      break;
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      auto* box = location_data_.mutable_relative_bounding_box();
      box->set_xmin(box->xmin() +
                    BestShift(box->xmin(), box->xmin() + box->width(), 1.0f));
      box->set_ymin(box->ymin() +
                    BestShift(box->ymin(), box->ymin() + box->height(), 1.0f));
      break;
    }
    case LocationData::MASK: {
      auto mask_bounding_box = MaskToRectangle(location_data_);
      const float x_shift = BestShift(mask_bounding_box.xmin(),
                                      mask_bounding_box.xmax(), image_width);
      const float y_shift = BestShift(mask_bounding_box.xmin(),
                                      mask_bounding_box.xmax(), image_height);
      auto* mask = location_data_.mutable_mask();
      CHECK_EQ(image_width, mask->width());
      CHECK_EQ(image_height, mask->height());
      for (auto& interval :
           *mask->mutable_rasterization()->mutable_interval()) {
        interval.set_y(interval.y() + y_shift);
        interval.set_left_x(interval.left_x() + x_shift);
        interval.set_right_x(interval.right_x() + x_shift);
      }
      break;
    }
  }
  return *this;
}

Location& Location::Crop(const Rectangle_i& crop_box) {
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      // Do nothing.
      break;
    }
    case LocationData::BOUNDING_BOX: {
      auto* box = location_data_.mutable_bounding_box();
      const int xmin = std::max(box->xmin(), crop_box.xmin());
      const int ymin = std::max(box->ymin(), crop_box.ymin());
      const int xmax = std::min(box->width() + box->xmin(), crop_box.xmax());
      const int ymax = std::min(box->height() + box->ymin(), crop_box.ymax());
      box->set_xmin(xmin - crop_box.xmin());
      box->set_ymin(ymin - crop_box.ymin());
      box->set_width(xmax - xmin);
      box->set_height(ymax - ymin);
      break;
    }
    case LocationData::RELATIVE_BOUNDING_BOX:
      LOG(FATAL)
          << "Can't crop a relative bounding box using absolute coordinates. "
             "Use the 'Rectangle_f version of Crop() instead";
    case LocationData::MASK: {
      LocationData::BinaryMask new_mask;
      new_mask.set_width(crop_box.Width());
      new_mask.set_height(crop_box.Height());
      auto* rasterization = new_mask.mutable_rasterization();
      for (const auto& interval :
           location_data_.mask().rasterization().interval()) {
        if (interval.y() >= crop_box.ymin() && interval.y() < crop_box.ymax() &&
            interval.left_x() < crop_box.xmax() &&
            interval.right_x() > crop_box.xmin()) {
          auto* new_interval = rasterization->add_interval();
          new_interval->set_y(interval.y() - crop_box.ymin());
          new_interval->set_left_x(
              std::max(interval.left_x() - crop_box.xmin(), 0));
          new_interval->set_right_x(
              std::min(interval.right_x() - crop_box.xmin(), crop_box.Width()));
        }
      }
      location_data_.mutable_mask()->Swap(&new_mask);
      break;
    }
  }
  return *this;
}

Location& Location::Crop(const Rectangle_f& crop_box) {
  switch (location_data_.format()) {
    case LocationData::GLOBAL:
      // Do nothing.
      break;
    case LocationData::BOUNDING_BOX:
      LOG(FATAL)
          << "Can't crop an absolute bounding box using relative coordinates. "
             "Use the 'Rectangle_i version of Crop() instead";
    case LocationData::RELATIVE_BOUNDING_BOX: {
      auto* box = location_data_.mutable_relative_bounding_box();
      float right = box->xmin() + box->width();
      float bottom = box->ymin() + box->height();
      box->set_xmin(std::max(crop_box.xmin(), box->xmin()));
      box->set_ymin(std::max(crop_box.ymin(), box->ymin()));
      float newRight = std::min(crop_box.xmax(), right);
      float newBottom = std::min(crop_box.ymax(), bottom);
      box->set_width(newRight - box->xmin());
      box->set_height(newBottom - box->ymin());
      break;
    }
    case LocationData::MASK:
      LOG(FATAL) << "Can't crop a mask using relative coordinates. Use the "
                    "'Rectangle_i' version of Crop() instead";
  }
  return *this;
}

template <>
Rectangle_i Location::ConvertToBBox<Rectangle_i>(int image_width,
                                                 int image_height) const {
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      return Rectangle_i(0, 0, image_width, image_height);
    }
    case LocationData::BOUNDING_BOX: {
      const auto& box = location_data_.bounding_box();
      return Rectangle_i(box.xmin(), box.ymin(), box.width(), box.height());
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      const auto& box = location_data_.relative_bounding_box();
      // Taking the floor rather than rounding for the width and height ensures
      // that if the original relative bounding box was within the image bounds,
      // the absolute bounding box that it is converted to will also be
      // within the image bounds.
      return Rectangle_i(
          static_cast<int>(std::round(image_width * box.xmin())),
          static_cast<int>(std::round(image_height * box.ymin())),
          static_cast<int>(image_width * box.width()),
          static_cast<int>(image_height * box.height()));
    }
    case LocationData::MASK: {
      return MaskToRectangle(location_data_);
    }
    default: {
      // Note that this shouldn't happen, but a default clause is required by
      // the Android compiler.
      return Rectangle_i();
    }
  }
}

Rectangle_f Location::GetRelativeBBox() const {
  CHECK_EQ(LocationData::RELATIVE_BOUNDING_BOX, location_data_.format());
  const auto& box = location_data_.relative_bounding_box();
  return Rectangle_f(box.xmin(), box.ymin(), box.width(), box.height());
}

Rectangle_f Location::ConvertToRelativeBBox(int image_width,
                                            int image_height) const {
  switch (location_data_.format()) {
    case LocationData::GLOBAL: {
      return Rectangle_f(0.0f, 0.0f, 1.0f, 1.0f);
    }
    case LocationData::BOUNDING_BOX: {
      const auto& box = location_data_.bounding_box();
      return Rectangle_f(static_cast<float>(box.xmin()) / image_width,
                         static_cast<float>(box.ymin()) / image_height,
                         static_cast<float>(box.width()) / image_width,
                         static_cast<float>(box.height()) / image_height);
    }
    case LocationData::RELATIVE_BOUNDING_BOX: {
      const auto& box = location_data_.relative_bounding_box();
      return Rectangle_f(box.xmin(), box.ymin(), box.width(), box.height());
    }
    case LocationData::MASK: {
      auto rect = MaskToRectangle(location_data_);
      return Rectangle_f(static_cast<float>(rect.xmin()) / image_width,
                         static_cast<float>(rect.ymin()) / image_height,
                         static_cast<float>(rect.Width()) / image_width,
                         static_cast<float>(rect.Height()) / image_height);
    }
    default: {
      // Note that this shouldn't happen, but a default clause is required by
      // the Android compiler.
      return Rectangle_f();
    }
  }
}

template <>
::mediapipe::BoundingBox Location::GetBBox<::mediapipe::BoundingBox>() const {
  CHECK_EQ(LocationData::BOUNDING_BOX, location_data_.format());
  const auto& box = location_data_.bounding_box();
  ::mediapipe::BoundingBox bounding_box;
  bounding_box.set_left_x(box.xmin());
  bounding_box.set_upper_y(box.ymin());
  bounding_box.set_right_x(box.width() + box.xmin());
  bounding_box.set_lower_y(box.height() + box.ymin());
  return bounding_box;
}

template <>
::mediapipe::BoundingBox Location::ConvertToBBox<::mediapipe::BoundingBox>(
    int image_width, int image_height) const {
  const auto& rect = ConvertToBBox<Rectangle_i>(image_width, image_height);
  ::mediapipe::BoundingBox bounding_box;
  bounding_box.set_left_x(rect.xmin());
  bounding_box.set_upper_y(rect.ymin());
  bounding_box.set_right_x(rect.xmax());
  bounding_box.set_lower_y(rect.ymax());
  return bounding_box;
}

std::vector<Point2_f> Location::GetRelativeKeypoints() const {
  CHECK_EQ(LocationData::RELATIVE_BOUNDING_BOX, location_data_.format());
  std::vector<Point2_f> keypoints;
  for (const auto& keypoint : location_data_.relative_keypoints()) {
    keypoints.emplace_back(Point2_f(keypoint.x(), keypoint.y()));
  }
  return keypoints;
}

std::vector<Point2_i> Location::ConvertToKeypoints(int image_width,
                                                   int image_height) const {
  std::vector<Point2_i> keypoints;
  for (const auto& keypoint : location_data_.relative_keypoints()) {
    keypoints.emplace_back(
        Point2_i(static_cast<int>(std::round(image_width * keypoint.x())),
                 static_cast<int>(std::round(image_height * keypoint.y()))));
  }
  return keypoints;
}

void Location::SetRelativeKeypoints(const std::vector<Point2_f>& keypoints) {
  location_data_.clear_relative_keypoints();
  for (const auto& keypoint : keypoints) {
    auto* relative_keypoint = location_data_.add_relative_keypoints();
    relative_keypoint->set_x(keypoint.x());
    relative_keypoint->set_y(keypoint.y());
  }
}

void Location::SetFromProto(const LocationData& proto) {
  location_data_ = proto;
}

void Location::ConvertToProto(LocationData* proto) const {
  *proto = location_data_;
}

LocationData Location::ConvertToProto() const {
  LocationData location_data;
  ConvertToProto(&location_data);
  return location_data;
}

}  // namespace mediapipe