///////////////////////////////////////////////////////////////////////

// File:        stridemap.cpp

// Description: Indexing into a 4-d tensor held in a 2-d Array.

// Author:      Ray Smith

//

// (C) Copyright 2016, Google Inc.

// 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 "stridemap.h"

#include <cassert> // for assert

namespace tesseract {

// Returns true if *this is a valid index.

bool StrideMap::Index::IsValid() const {

  // Cheap check first.

  for (int index : indices_) {

    if (index < 0) {

      return false;

    }

  }

  for (int d = 0; d < FD_DIMSIZE; ++d) {

    if (indices_[d] > MaxIndexOfDim(static_cast<FlexDimensions>(d))) {

      return false;

    }

  }

  return true;

}

// Returns true if the index of the given dimension is the last.

bool StrideMap::Index::IsLast(FlexDimensions dimension) const {

  return MaxIndexOfDim(dimension) == indices_[dimension];

}

// Given that the dimensions up to and including dim-1 are valid, returns the

// maximum index for dimension dim.

int StrideMap::Index::MaxIndexOfDim(FlexDimensions dim) const {

  int max_index = stride_map_->shape_[dim] - 1;

  if (dim == FD_BATCH) {

    return max_index;

  }

  assert(0 <= indices_[FD_BATCH]);

  const size_t batch = indices_[FD_BATCH];

  if (dim == FD_HEIGHT) {

    if (batch >= stride_map_->heights_.size() || stride_map_->heights_[batch] > max_index) {

      return max_index;

    }

    return stride_map_->heights_[batch] - 1;

  }

  if (batch >= stride_map_->widths_.size() || stride_map_->widths_[batch] > max_index) {

    return max_index;

  }

  return stride_map_->widths_[batch] - 1;

}

// Adds the given offset to the given dimension. Returns true if the result

// makes a valid index.

bool StrideMap::Index::AddOffset(int offset, FlexDimensions dimension) {

  indices_[dimension] += offset;

  SetTFromIndices();

  return IsValid();

}

// Increments the index in some encapsulated way that guarantees to remain

// valid until it returns false, meaning that the iteration is complete.

bool StrideMap::Index::Increment() {

  for (int d = FD_DIMSIZE - 1; d >= 0; --d) {

    if (!IsLast(static_cast<FlexDimensions>(d))) {

      t_ += stride_map_->t_increments_[d];

      ++indices_[d];

      return true;

    }

    t_ -= stride_map_->t_increments_[d] * indices_[d];

    indices_[d] = 0;

    // Now carry to the next dimension.

  }

  return false;

}

// Decrements the index in some encapsulated way that guarantees to remain

// valid until it returns false, meaning that the iteration (that started

// with InitToLast()) is complete.

bool StrideMap::Index::Decrement() {

  for (int d = FD_DIMSIZE - 1; d >= 0; --d) {

    if (indices_[d] > 0) {

      --indices_[d];

      if (d == FD_BATCH) {

        // The upper limits of the other dimensions may have changed as a result

        // of a different batch index, so they have to be reset.

        InitToLastOfBatch(indices_[FD_BATCH]);

      } else {

        t_ -= stride_map_->t_increments_[d];

      }

      return true;

    }

    indices_[d] = MaxIndexOfDim(static_cast<FlexDimensions>(d));

    t_ += stride_map_->t_increments_[d] * indices_[d];

    // Now borrow from the next dimension.

  }

  return false;

}

// Initializes the indices to the last valid location in the given batch

// index.

void StrideMap::Index::InitToLastOfBatch(int batch) {

  indices_[FD_BATCH] = batch;

  for (int d = FD_BATCH + 1; d < FD_DIMSIZE; ++d) {

    indices_[d] = MaxIndexOfDim(static_cast<FlexDimensions>(d));

  }

  SetTFromIndices();

}

// Computes and sets t_ from the current indices_.

void StrideMap::Index::SetTFromIndices() {

  t_ = 0;

  for (int d = 0; d < FD_DIMSIZE; ++d) {

    t_ += stride_map_->t_increments_[d] * indices_[d];

  }

}

// Sets up the stride for the given array of height, width pairs.

void StrideMap::SetStride(const std::vector<std::pair<int, int>> &h_w_pairs) {

  int max_height = 0;

  int max_width = 0;

  for (const std::pair<int, int> &hw : h_w_pairs) {

    int height = hw.first;

    int width = hw.second;

    heights_.push_back(height);

    widths_.push_back(width);

    if (height > max_height) {

      max_height = height;

    }

    if (width > max_width) {

      max_width = width;

    }

  }

  shape_[FD_BATCH] = heights_.size();

  shape_[FD_HEIGHT] = max_height;

  shape_[FD_WIDTH] = max_width;

  ComputeTIncrements();

}

// Scales width and height dimensions by the given factors.

void StrideMap::ScaleXY(int x_factor, int y_factor) {

  for (int &height : heights_) {

    height /= y_factor;

  }

  for (int &width : widths_) {

    width /= x_factor;

  }

  shape_[FD_HEIGHT] /= y_factor;

  shape_[FD_WIDTH] /= x_factor;

  ComputeTIncrements();

}

// Reduces width to 1, across the batch, whatever the input size.

void StrideMap::ReduceWidthTo1() {

  widths_.assign(widths_.size(), 1);

  shape_[FD_WIDTH] = 1;

  ComputeTIncrements();

}

// Transposes the width and height dimensions.

void StrideMap::TransposeXY() {

  std::swap(shape_[FD_HEIGHT], shape_[FD_WIDTH]);

  std::swap(heights_, widths_);

  ComputeTIncrements();

}

// Computes t_increments_ from shape_.

void StrideMap::ComputeTIncrements() {

  t_increments_[FD_DIMSIZE - 1] = 1;

  for (int d = FD_DIMSIZE - 2; d >= 0; --d) {

    t_increments_[d] = t_increments_[d + 1] * shape_[d + 1];

  }

}

} // namespace tesseract