/**
 * Copyright 2020 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 <vector>
#include <algorithm>
#include "minddata/dataset/engine/opt/pre/epoch_injection_pass.h"
#include "minddata/dataset/engine/execution_tree.h"
#include "minddata/dataset/engine/datasetops/epoch_ctrl_op.h"
#include "minddata/dataset/engine/datasetops/device_queue_op.h"

namespace mindspore {
namespace dataset {

// constructor
EpochInjectionPass::InjectionFinder::InjectionFinder(std::shared_ptr<DatasetOp> node) : injection_point_(node) {}

// Performs finder work for BuildVocabOp that has special rules about epoch control injection
Status EpochInjectionPass::InjectionFinder::PreRunOnNode(std::shared_ptr<BuildVocabOp> node, bool *modified) {
  injection_point_ = nullptr;
  return Status::OK();
}

// Performs finder work for BuildSentencePieceVocabOp that has special rules about epoch control injection
Status EpochInjectionPass::InjectionFinder::PreRunOnNode(std::shared_ptr<BuildSentencePieceVocabOp> node,
                                                         bool *modified) {
  injection_point_ = nullptr;
  return Status::OK();
}

// Temporary code to prevent the injection of epoch control when cache op is present
// Remove this code in cache op phase 2
Status EpochInjectionPass::InjectionFinder::PreRunOnNode(std::shared_ptr<CacheOp> node, bool *modified) {
  injection_point_ = nullptr;
  return Status::OK();
}

Status EpochInjectionPass::InjectionFinder::RunOnNode(std::shared_ptr<DeviceQueueOp> node, bool *modified) {
  // Assumption: There is only one DeviceQueueOp in a pipeline. This assumption is not validated here.
  injection_point_ = node->child(0);
  return Status::OK();
}

// constructor
EpochInjectionPass::EpochInjectionPass() {}

// Runs an injection pass to inject in operators needed at the pre pass stage
Status EpochInjectionPass::RunOnTree(ExecutionTree *tree, bool *modified) {
  MS_LOG(INFO) << "Pre pass: Injection pass started.";

  // First, run the finder to perform any injection info before we can go ahead to drive the op injection work.
  // The finder can make updates to the EpochInjectionPass object.
  EpochInjectionPass::InjectionFinder finder(tree->root());
  RETURN_IF_NOT_OK(finder.Run(tree, modified));

  // The first injection logic is to check if we should inject the epoch control op as the root node.
  // Do not inject the op if the number of epochs is 1.
  int32_t num_epochs = tree->num_epochs();
  std::shared_ptr<DatasetOp> epoch_inject_node = finder.injection_point();
  if (num_epochs != 1 && epoch_inject_node != nullptr) {
    std::shared_ptr<EpochCtrlOp> epoch_ctrl_op;
    RETURN_IF_NOT_OK(EpochCtrlOp::Builder(num_epochs).Build(&epoch_ctrl_op));
    RETURN_IF_NOT_OK(tree->AssociateNode(epoch_ctrl_op));
    epoch_inject_node->InsertAsParent(epoch_ctrl_op);
  }

  MS_LOG(INFO) << "Pre pass: Injection pass complete.";
  return Status::OK();
}
}  // namespace dataset
}  // namespace mindspore
