/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 "singa/io/transformer.h"
#include <time.h>

#ifdef USE_OPENCV
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#endif

namespace singa {

  Tensor ImageTransformer::Apply(int flag, Tensor& input) {
    CHECK_LE(input.nDim(), 4u);
    CHECK_GE(input.nDim(), 2u);
    CHECK_EQ(input.data_type(), kFloat32) << "Data type " << input.data_type()
      << " is invalid for an raw image";
    srand((unsigned int)time(NULL));
    /// TODO
    /// currently only consider one sample each time

    /// resize image using opencv resize
    Tensor temp1;
#ifdef USE_OPENCV
    temp1 = resize(input, resize_height_, resize_width_, image_dim_order_);
#else
    temp1 = input;
#endif

    /// crop
    Tensor temp2;
    size_t height = 0, width = 0;
    if (input.nDim() >= 3u) {
      if (image_dim_order_ == "CHW")
        height = temp1.shape(input.nDim() - 2), width = temp1.shape(input.nDim() - 1);
      else if (image_dim_order_ == "HWC")
        height = temp1.shape(input.nDim() - 3), width = temp1.shape(input.nDim() - 2);
      else
        LOG(FATAL) << "Unknow dimension order for images " << image_dim_order_
               << " Only support 'HWC' and 'CHW'";
    } else /// input is 2D gray image
      height = temp1.shape(0), width = temp1.shape(1);

    if (crop_shape_.size() == 2) {
      if (flag == kTrain) { 
        /// random crop
        if (crop_shape_[0] > height || crop_shape_[0] > width)
          LOG(FATAL) << "Crop size larger than the size of raw image";
        size_t crop_h_offset = rand() % ((height - crop_shape_[0]) / 2), 
               crop_w_offset = rand() % ((width - crop_shape_[1]) / 2);
        temp2 = crop(temp1, crop_shape_[0], crop_shape_[1], 
                  crop_h_offset, crop_w_offset, image_dim_order_);
      } else if (flag == kEval) {
        /// central crop
        size_t crop_h_offset = (height - crop_shape_[0]) / 2,
               crop_w_offset = (width - crop_shape_[1]) / 2;
        temp2 = crop(temp1, crop_shape_[0], crop_shape_[1], 
                  crop_h_offset, crop_w_offset, image_dim_order_); 
      }
    }
    else temp2 = temp1;

    /// mirror
    Tensor output;
    if ((flag == kTrain) && (rand() % 2))
        output = mirror(temp2, true, false, image_dim_order_);
    else output = temp2;
    return output;
  }

#ifdef USE_OPENCV
  Tensor resize(Tensor& input, const size_t resize_height, 
               const size_t resize_width, const string& image_dim_order) {
    CHECK_LE(input.nDim(), 4u);
    CHECK_GE(input.nDim(), 2u);
    if (!resize_height || !resize_width) return input;
    Tensor output;
    cv::Mat mat;
    const auto* in = input.data<float>();
    if (input.nDim() == 4u) {
      /// TODO
      /// batch based resize
      LOG(FATAL) << "Not implemented";
    } else if (input.nDim() == 3u) {
      if (image_dim_order == "CHW") {
        size_t height = input.shape(1), width = input.shape(2),
               channel = input.shape(0);
        if (channel == 3u) {
          mat = cv::Mat(height, width, CV_32FC3, cv::Scalar(0, 0, 0));
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              for (size_t k = 0; k < channel; k++)
                mat.at<cv::Vec3f>(i, j)[k] = in[k * height * width + i * width + j];
        }
        else if (channel == 1u) {
          mat = cv::Mat(height, width, CV_32FC1);
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
                mat.at<cv::Vec<float, 1>>(i, j)[0] = in[i * width + j];
        }
        else LOG(FATAL) << "Invalid channel size: " << channel;
      } else if (image_dim_order == "HWC") {
        size_t height = input.shape(0), width = input.shape(1),
               channel = input.shape(2);
        if (channel == 3u) {
          mat = cv::Mat(height, width, CV_32FC3, cv::Scalar(0, 0, 0));
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              for (size_t k = 0; k < channel; k++)
                mat.at<cv::Vec3f>(i, j)[k] =
                  in[i * width * channel + j * channel + k];
        } else if (channel == 1u) { /// 2D gray image
          mat = cv::Mat(height, width, CV_32FC1);
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              mat.at<cv::Vec<float, 1>>(i, j)[0] = in[i * width + j];
        } else LOG(FATAL) << "Invalid channel size: " << channel;
      } else {
        LOG(FATAL) << "Unknow dimension order for images " << image_dim_order
                   << " Only support 'HWC' and 'CHW'";
      }
    } else { /// 2D gray image
      size_t height = input.shape(0), width = input.shape(1);
      mat = cv::Mat(height, width, CV_32FC1);
      for (size_t i = 0; i < height; i++)
        for (size_t j = 0; j < width; j++)
          mat.at<cv::Vec<float, 1>>(i, j)[0] = in[i * width + j];
    }
    cv::Size size(resize_width, resize_height);
    cv::Mat resized;
    cv::resize(mat, resized, size);
    CHECK_EQ(resized.size().height, resize_height);
    CHECK_EQ(resized.size().width, resize_width);
    size_t new_size = resize_height * resize_width * resized.channels();
    float* out = new float[new_size];
    if (input.nDim() == 4u) {
      /// TODO
      /// batch based resize
      LOG(FATAL) << "Not implemented";
    } else if (input.nDim() == 3u) {
      if (image_dim_order == "CHW") {
        size_t height = resize_height, width = resize_width,
           channel = input.shape(0);
        if (channel == 3u) {
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              for (size_t k = 0; k < channel; k++)
                out[k * height * width + i * width + j] = resized.at<cv::Vec3f>(i, j)[k];
        } else { /// 2D gray image
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              out[i * width + j] = resized.at<cv::Vec<float, 1>>(i, j)[0];
        }
        Tensor temp(Shape{channel, height, width});
        temp.CopyDataFromHostPtr<float>(out, new_size);
        output = temp;
      } else {
        size_t height = resize_height, width = resize_width,
           channel = input.shape(2);
        if (channel == 3u) {
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              for (size_t k = 0; k < channel; k++)
                out[i * width * channel + j * channel + k] = resized.at<cv::Vec3f>(i, j)[k];
        } else { /// 1 channel
          for (size_t i = 0; i < height; i++)
            for (size_t j = 0; j < width; j++)
              out[i * width + j] = resized.at<cv::Vec<float, 1>>(i, j)[0];
        }
        Tensor temp(Shape{height, width, channel}); 
        temp.CopyDataFromHostPtr<float>(out, new_size);
        output = temp;
      }
    } else { /// 2D gray image
      size_t height = resize_height, width = resize_width;
      for (size_t i = 0; i < height; i++)
        for (size_t j = 0; j < width; j++)
          out[i * width + j] = resized.at<cv::Vec<float, 1>>(i, j)[0];
      Tensor temp(Shape{height, width});
      temp.CopyDataFromHostPtr<float>(out, new_size);
      output = temp;
    }
    delete[] out;
    return output;
  }
#endif

  Tensor crop(Tensor& input, const size_t crop_height, const size_t crop_width, 
             const size_t crop_h_offset, const size_t crop_w_offset, 
             const string& image_dim_order) {
    CHECK_LE(input.nDim(), 4u);
    CHECK_GE(input.nDim(), 2u);

    Tensor output;
    const float* in = input.data<float>();
    size_t out_idx = 0, in_idx = 0;
    if (input.nDim() == 4u) {
      /// TODO
      LOG(FATAL) << "Not implemented";
    } else if (input.nDim() == 3u) {
      if (image_dim_order == "CHW") {
        size_t height = input.shape(1), width = input.shape(2),
            channel = input.shape(0); 
        CHECK_LE(crop_height + crop_h_offset, height);
        CHECK_LE(crop_width + crop_w_offset, width);
        float* out = new float[crop_height * crop_width * channel];
        for (size_t c = 0; c < channel; c++) {
          for (size_t h = 0; h < crop_height; h++) {
            for (size_t w = 0; w < crop_width; w++) {
              in_idx = (c * height + crop_h_offset + h) * width + crop_w_offset + w;
              out_idx = (c * crop_height + h) * crop_width + w;
              out[out_idx] = in[in_idx];
            }
          }
        }
        output.Reshape(Shape{channel, crop_height, crop_width});
        output.CopyDataFromHostPtr<float>(out, crop_height * crop_width * channel);
        delete[] out;
      } else if (image_dim_order == "HWC") {
        size_t height = input.shape(0), width = input.shape(1), 
               channel = input.shape(2); 
        CHECK_LE(crop_height + crop_h_offset, height);
        CHECK_LE(crop_width + crop_w_offset, width);
        float* out = new float[crop_height * crop_width * channel];
        for (size_t c = 0; c < channel; c++) {
          for (size_t h = 0; h < crop_height; h++) {
            for (size_t w = 0; w < crop_width; w++) {
              in_idx = ((crop_h_offset + h) * width + crop_w_offset + w) * channel + c;
              out_idx = (h * crop_width + w) * channel + c;
              out[out_idx] = in[in_idx];
            }
          }
        }
        output.Reshape(Shape{crop_height, crop_width, channel});
        output.CopyDataFromHostPtr<float>(out, crop_height * crop_width * channel);
        delete[] out;
      } else {
        LOG(FATAL) << "Unknow dimension order for images " << image_dim_order
                   << " Only support 'HWC' and 'CHW'";
      }
    } else { /// 2D gray image
      size_t height = input.shape(0), width = input.shape(1); 
      CHECK_LE(crop_height + crop_h_offset, height);
      CHECK_LE(crop_width + crop_w_offset, width);
      float* out = new float[crop_height * crop_width];
      for (size_t h = 0; h < crop_height; h++) {
        for (size_t w = 0; w < crop_width; w++) {
          in_idx = (crop_h_offset + h) * width + crop_w_offset + w;
          out_idx = h * crop_width + w;
          out[out_idx] = in[in_idx];
        }
      }
      output.Reshape(Shape{crop_height, crop_width});
      output.CopyDataFromHostPtr<float>(out, crop_height * crop_width);
      delete[] out;
    }
    return output;
  }

  Tensor mirror(Tensor& input, const bool horizontal_mirror,
             const bool vertical_mirror, const string& image_dim_order) {
    CHECK_LE(input.nDim(), 4u);
    CHECK_GE(input.nDim(), 2u);
    if (!horizontal_mirror && !vertical_mirror) return input;

    Tensor output;
    const float* in = input.data<float>();
    size_t out_idx = 0, in_idx = 0;
    if (input.nDim() == 4u) {
      /// TODO
      LOG(FATAL) << "Not implemented";
    } else if (input.nDim() == 3u) {
      if (image_dim_order == "CHW") {
        size_t height = input.shape(1), width = input.shape(2),
            channel = input.shape(0);
        float* out = new float[height * width * channel];
        for (size_t c = 0; c < channel; c++) {
          for (size_t h = 0; h < height; h++) {
            for (size_t w = 0; w < width; w++) {
              in_idx = (c * height + h) * width + w;
              if (horizontal_mirror && vertical_mirror)
                out_idx = (c * height + (height - 1 - h)) * width + (width - 1 - w);
              else if (horizontal_mirror)
                out_idx = (c * height + h) * width + (width - 1 - w);
              else /// only do vertical mirror
                out_idx = (c * height + (height - 1 - h)) * width + w;
              out[out_idx] = in[in_idx];
            }
          }
        }
        output.Reshape(Shape{channel, height, width});
        output.CopyDataFromHostPtr<float>(out, height * width * channel);
        delete[] out;
      } else if (image_dim_order == "HWC") {
        size_t height = input.shape(0), width = input.shape(1),
            channel = input.shape(2);
        float* out = new float[height * width * channel];
        for (size_t c = 0; c < channel; c++) {
          for (size_t h = 0; h < height; h++) {
            for (size_t w = 0; w < width; w++) {
              in_idx = (h * width + w) * channel + c;
              if (horizontal_mirror && vertical_mirror)
                out_idx = ((height - 1 - h) * width + (width - 1 - w)) * channel + c;
              else if (horizontal_mirror)
                out_idx = (h * width + (width - 1 - w)) * channel + c;
              else /// only do vertical mirror
                out_idx = ((height - 1 - h) * width + w) * channel + c;
              out[out_idx] = in[in_idx];
            }
          }
        }
        output.Reshape(Shape{height, width, channel});
        output.CopyDataFromHostPtr<float>(out, height * width * channel);
        delete[] out;
      } else {
        LOG(FATAL) << "Unknow dimension order for images " << image_dim_order
                   << " Only support 'HWC' and 'CHW'";
      }
    } else { /// 2D gray image
      size_t height = input.shape(0), width = input.shape(1);
      float* out = new float[height * width];
      for (size_t h = 0; h < height; h++) {
        for (size_t w = 0; w < width; w++) {
          in_idx = h * width + w;
          if (horizontal_mirror && vertical_mirror)
            out_idx = (height - 1 - h) * width + (width - 1 - w);
          else if (horizontal_mirror)
            out_idx = h * width + (width - 1 - w);
          else /// only do vertical mirror
            out_idx = (height - 1 - h) * width + w;
          out[out_idx] = in[in_idx];
        }
      }
      output.Reshape(Shape{height, width});
      output.CopyDataFromHostPtr<float>(out, height * width);
      delete[] out;
    }
    return output;
  }
} // namespace singa