/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * 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 "dataset/engine/dataset_iterator.h"
#include <unordered_map>
#include <utility>
#include "dataset/core/data_type.h"
#include "dataset/core/tensor.h"
#include "dataset/core/tensor_shape.h"
#include "dataset/engine/data_buffer.h"
#include "dataset/engine/execution_tree.h"
#include "dataset/util/status.h"
#include "dataset/engine/datasetops/dataset_op.h"

namespace mindspore {
namespace dataset {
// Constructor of the IteratorBase
IteratorBase::IteratorBase() : curr_buffer_(nullptr), eof_handled_(false) {}

IteratorBase::~IteratorBase() = default;

// Fetches one row of data from the iterator as a column map.
Status IteratorBase::GetNextAsMap(TensorMap *out_map) {
  if (out_map == nullptr) {
    RETURN_STATUS_UNEXPECTED("Null output map in iterator!");
  }

  out_map->clear();

  TensorRow curr_row;
  RETURN_IF_NOT_OK(FetchNextTensorRow(&curr_row));

  // Return empty map if there's no data
  if (curr_row.empty()) {
    return Status::OK();
  }

  // The column name mapping is needed to be able to produce the tensor map output.
  // The column name mapping comes from the source operator that is producing the data into the iterator.
  // To avoid having to fetch this for every time, we'll take a local copy of the column name id mapping
  // and save in the iterator.  We only have to do this once.  All subsequent iterations use the same mapping.
  if (col_name_id_map_.empty()) {
    // Determine the column name map by calling the derived class method to retrieve the column
    // name map
    col_name_id_map_ = this->GetColumnNameMap();
  }

  // Populate the out map from the row and return it
  for (auto colMap : col_name_id_map_) {
    (*out_map)[colMap.first] = std::move(curr_row[colMap.second]);
  }

  return Status::OK();
}

// Fetches one row of data from the iterator.
// The base class version simply performs error handling and returns empty row. Actual
// functionality exists in the derived versions of this function.
Status IteratorBase::FetchNextTensorRow(TensorRow *out_row) {
  if (out_row == nullptr) {
    RETURN_STATUS_UNEXPECTED("Null output row in iterator!");
  }

  // clear the old tensor row
  out_row->clear();

  return Status::OK();
}

// Constructor of the DatasetIterator
DatasetIterator::DatasetIterator(std::shared_ptr<ExecutionTree> exe_tree)
    : IteratorBase(),
      root_(exe_tree->root()),
      tracing_(nullptr),
      cur_batch_num_(0),
      cur_connector_size_(0),
      cur_connector_capacity_(0) {
  std::shared_ptr<Tracing> node;
  Status s = exe_tree->GetProfilingManager()->GetTracingNode(kDatasetIteratorTracingName, &node);
  if (s.IsOk()) {
    tracing_ = std::dynamic_pointer_cast<DatasetIteratorTracing>(node);
  }
}

DatasetIterator::~DatasetIterator() = default;

// Fetches one row of data from the iterator.  Overrides the base class.  This one fetches
// from the tree root node directly.
Status DatasetIterator::FetchNextTensorRow(TensorRow *out_row) {
  // Common code init and error checking in the base class.
  RETURN_IF_NOT_OK(IteratorBase::FetchNextTensorRow(out_row));

  // Once eof is handled, always return empty row.  Class must be destroyed and recreated if you
  // want to iterate again.
  if (eof_handled_) {
    return Status::OK();
  }

  // Check if we need to get a new DataBuffer to iterate.
  if (curr_buffer_ == nullptr || curr_buffer_->NumRows() == 0) {
    if (tracing_ != nullptr) {
      cur_connector_size_ = root_->ConnectorSize();
      cur_connector_capacity_ = root_->ConnectorCapacity();
    }
    RETURN_IF_NOT_OK(root_->GetNextBuffer(&curr_buffer_));

    // Since GetNextBuffer was used rather than GetNextInput(), it means we need to manually
    // handle eoe and eof messages here.
    //
    // An eoe buffer means we have iterated fully to the end of the tree.
    // An eoe buffer will be immediately followed by an eof buffer, which signals the shutdown of
    // all operators.
    if (curr_buffer_->eoe()) {
      MS_LOG(DEBUG) << "End of data iteration. Fetch eof and then return empty row.";

      // Before returning the last empty vector, fetch the eof buffer which should be the last
      // buffer, and then free it.
      RETURN_IF_NOT_OK(root_->GetNextBuffer(&curr_buffer_));

      if (!curr_buffer_->eof()) {
        RETURN_STATUS_UNEXPECTED("Non-eof after getting eoe in iterator!");
      }
      eof_handled_ = true;
      curr_buffer_.reset();  // explicitly free the eof buffer
      // Set tree to Finished state
      root_->Tree()->SetFinished();

      return Status::OK();
    }

    if (curr_buffer_->eof()) {
      // An eof by itself, without being preceded by an eoe, is possible if a repeat operator
      // exists below us in the stack. Repeat operator eats eoe's but eventually allows the
      // flow of an eof up the pipeline by itself.
      eof_handled_ = true;
      curr_buffer_.reset();  // explicitly free the eof buffer
      // Set tree to Finished state
      root_->Tree()->SetFinished();
      return Status::OK();
    }
  }

  // If we got this far, now it's time to pop that next row for return to caller
  RETURN_IF_NOT_OK(curr_buffer_->PopRow(out_row));
  if (tracing_ != nullptr) {
    cur_batch_num_++;
    tracing_->Record(CONNECTOR_DEPTH, cur_connector_capacity_, cur_batch_num_, cur_connector_size_);
  }
  return Status::OK();
}

Status DatasetIterator::GetOutputShapes(std::vector<TensorShape> *out_shapes) {
  if (out_shapes == nullptr) {
    RETURN_STATUS_UNEXPECTED("Null output shape argument");
  }
  if (device_queue_row_.empty()) {
    RETURN_IF_NOT_OK(FetchNextTensorRow(&device_queue_row_));
  }
  for (auto ts : device_queue_row_) {
    out_shapes->push_back(ts->shape());
  }

  return Status::OK();
}

Status DatasetIterator::GetOutputTypes(std::vector<DataType> *out_types) {
  if (out_types == nullptr) {
    RETURN_STATUS_UNEXPECTED("Null output type argument");
  }
  if (device_queue_row_.empty()) {
    RETURN_IF_NOT_OK(FetchNextTensorRow(&device_queue_row_));
  }
  for (auto ts : device_queue_row_) {
    out_types->push_back(ts->type());
  }
  return Status::OK();
}

// Getter
std::unordered_map<std::string, int32_t> DatasetIterator::GetColumnNameMap() const {
  return root_->column_name_id_map();
}

// Constructor of the ChildIterator
ChildIterator::ChildIterator(DatasetOp *current_op, int32_t worker_id, int32_t child_idx)
    : IteratorBase(), current_op_(current_op), child_idx_(child_idx), worker_id_(worker_id), end_epoch_(false) {}

ChildIterator::~ChildIterator() { current_op_ = nullptr; }

// Fetches one row of data from the iterator.  Overrides the base class.  This one fetches
// only from the child/worker id as given from the constructor.
Status ChildIterator::FetchNextTensorRow(TensorRow *out_row) {
  // Common code init and error checking in the base class.
  RETURN_IF_NOT_OK(IteratorBase::FetchNextTensorRow(out_row));

  // Once eof is handled, always return empty row.  Class must be destroyed and recreated if you
  // want to iterate again.
  if (eof_handled_) {
    return Status::OK();
  }

  // Check if we need to get a new DataBuffer to iterate.
  if (curr_buffer_ == nullptr || curr_buffer_->NumRows() == 0) {
    RETURN_IF_NOT_OK(current_op_->GetNextInput(&curr_buffer_, worker_id_, child_idx_));

    // Unlike the DatasetIterator, this child iterator does not quit after eoe.
    // Instead, if an eoe is picked up here, we simply return an empty vector and it's up to the
    // caller to decide what it wants to do next.
    if (curr_buffer_->eoe()) {
      MS_LOG(DEBUG) << "Child iterator picked up EOE.";
      end_epoch_ = true;
      return Status::OK();
    }

    if (curr_buffer_->eof()) {
      MS_LOG(DEBUG) << "Child iterator picked up EOF.";
      eof_handled_ = true;
      return Status::OK();
    }
  }

  // If we got this far, now it's time to pop that next row for return to caller
  RETURN_IF_NOT_OK(curr_buffer_->PopRow(out_row));

  return Status::OK();
}

// drain till the next eoe
Status ChildIterator::Drain() {
  if (end_epoch_ == true) {
    // Calling drain against a child that is already at it's eoe state will not result in any action.
    // This allows you to do:
    // - fetch until empty row
    // - drain (will not actually drain because you are already at the end of the iteration)
    // However, the next time after that, it will perform it's normal draining activities.
    end_epoch_ = false;
    MS_LOG(DEBUG) << "No operation drain, already at end of epoch.";
    return Status::OK();
  }
  MS_LOG(DEBUG) << "Child draining buffers until eoe.";
  // else we drain until eoe or eof, eof here is for sanity check
  while (!curr_buffer_->eoe() && !curr_buffer_->eof()) {
    RETURN_IF_NOT_OK(current_op_->GetNextInput(&curr_buffer_, worker_id_, child_idx_));
  }
  if (curr_buffer_->eof()) {
    return Status(StatusCode::kUnexpectedError, __LINE__, __FILE__, "Child iterator picked up EOF in drain.");
  }
  return Status::OK();
}

// Getter
std::unordered_map<std::string, int32_t> ChildIterator::GetColumnNameMap() const {
  return current_op_->child(child_idx_)->column_name_id_map();
}
}  // namespace dataset
}  // namespace mindspore