////////////////////////////////////////////////////////////////////////////////
// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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 "lbann/transforms/vision/center_crop.hpp"
#include "lbann/utils/dim_helpers.hpp"
#include "lbann/utils/memory.hpp"
#include "lbann/utils/opencv.hpp"

#include "lbann/proto/transforms.pb.h"

#include <cmath>

namespace lbann {
namespace transform {

void center_crop::apply(utils::type_erased_matrix& data,
                        std::vector<size_t>& dims)
{
  cv::Mat src = utils::get_opencv_mat(data, dims);
  if (dims[1] <= m_h || dims[2] <= m_w) {
    std::stringstream ss;
    ss << "Center crop to " << m_h << "x" << m_w << " applied to input "
       << dims[1] << "x" << dims[2];
    LBANN_ERROR(ss.str());
  }
  std::vector<size_t> new_dims = {dims[0], m_h, m_w};
  auto dst_real = El::Matrix<uint8_t>(get_linear_size(new_dims), 1);
  cv::Mat dst = utils::get_opencv_mat(dst_real, new_dims);
  // Compute upper-left corner of crop.
  const size_t x = std::round(float(src.cols - m_w) / 2.0);
  const size_t y = std::round(float(src.rows - m_h) / 2.0);
  // Sanity check.
  if (x >= static_cast<size_t>(src.cols) ||
      y >= static_cast<size_t>(src.rows) ||
      (x + m_w) > static_cast<size_t>(src.cols) ||
      (y + m_h) > static_cast<size_t>(src.rows)) {
    std::stringstream ss;
    ss << "Bad crop dimensions for " << src.rows << "x" << src.cols << ": "
       << m_h << "x" << m_w << " at (" << x << "," << y << ")";
    LBANN_ERROR(ss.str());
  }
  // Copy is needed to ensure this is continuous.
  src(cv::Rect(x, y, m_h, m_w)).copyTo(dst);
  data.emplace<uint8_t>(std::move(dst_real));
  dims = new_dims;
}

std::unique_ptr<transform>
build_center_crop_transform_from_pbuf(google::protobuf::Message const& msg)
{
  auto const& params =
    dynamic_cast<lbann_data::Transform::CenterCrop const&>(msg);
  return std::make_unique<center_crop>(params.height(), params.width());
}

} // namespace transform
} // namespace lbann
