// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// 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 "paddle/fluid/framework/new_executor/interpreter/plan.h"

#include "paddle/fluid/framework/program_desc.h"

namespace paddle {
namespace framework {
namespace interpreter {

Plan::Plan(const std::vector<std::shared_ptr<Job>>& job_list,
           const std::unordered_map<std::string, ProgramDesc*>& type_to_program)
    : job_list_(job_list),
      type_to_program_(type_to_program),
      micro_batch_num_(1) {
  for (size_t i = 0; i < job_list_.size(); ++i) {
    const auto& job = job_list_[i];
    PADDLE_ENFORCE(type_to_program_.find(job->Type()) != type_to_program_.end(),
                   phi::errors::InvalidArgument(
                       "The %d-th job (type:%s, micro_batch_id:%d) has no "
                       "corresponding Program.",
                       i,
                       job->Type(),
                       job->MicroBatchId()));

    micro_batch_num_ = std::max(micro_batch_num_, job->MicroBatchId() + 1);
  }
}

const std::vector<std::shared_ptr<Job>>& Plan::JobList() const {
  return job_list_;
}

const ProgramDesc* Plan::Program(const std::string& job_type) const {
  return type_to_program_.at(job_type);
}

int64_t Plan::MicroBatchNum() const { return micro_batch_num_; }

}  // namespace interpreter
}  // namespace framework
}  // namespace paddle