#include "dlr_pipeline.h"

#include <stdlib.h>

#include <cstring>
#include <fstream>
#include <iterator>
#include <numeric>

using namespace dlr;

void PipelineModel::CheckModelsCompatibility(const DLRModelPtr& m0, const DLRModelPtr& m1,
                                             const int m1_id, const bool is_runtime_check) {
  if (!is_runtime_check) {
    CHECK_EQ(m0->GetNumOutputs(), m1->GetNumInputs())
        << "Number of outputs/inputs mismatch between models" << m1_id - 1 << " and " << m1_id
        << std::endl;
  }
  // Check each model input
  for (int j = 0; j < m1->GetNumInputs(); j++) {
    int64_t out_size;
    int out_dim;
    m0->GetOutputSizeDim(j, &out_size, &out_dim);
    if (!is_runtime_check) {
      CHECK_EQ(out_dim, m1->GetInputDim(j))
          << "Number of dimensions mismatch between output/input #" << j << ", models " << m1_id - 1
          << " and " << m1_id << std::endl;
    }
    const int64_t in_size = m1->GetInputSize(j);
    // Skip the check for dynamic sizes (negative size)
    if (out_size >= 0 && in_size >= 0) {
      CHECK_EQ(out_size, m1->GetInputSize(j))
          << "Size mismatch between output/input #" << j << ", models " << m1_id - 1 << " and "
          << m1_id << std::endl;
    }
    if (!is_runtime_check) {
      CHECK_EQ(strcmp(m0->GetOutputType(j), m1->GetInputType(j)), 0)
          << "Type mismatch between output/input #" << j << ", models " << m1_id - 1 << " and "
          << m1_id << std::endl;
    }
    std::vector<int64_t> in_shape = m1->GetInputShape(j);
    std::vector<int64_t> out_shape(in_shape.size());
    m0->GetOutputShape(j, out_shape.data());
    for (int k = 0; k < in_shape.size(); k++) {
      int64_t in_shape_elem = in_shape[k];
      int64_t out_shape_elem = out_shape[k];
      // Skip the check for dynamic shape elements (-1)
      if (in_shape_elem >= 0 && out_shape_elem >= 0) {
        CHECK_EQ(in_shape_elem, out_shape_elem)
            << "Shape mismatch between output/input #" << j << ", models " << m1_id - 1 << " and "
            << m1_id << std::endl;
      }
    }
  }
}

void PipelineModel::SetupPipelineModel() {
  CHECK_GT(dlr_models_.size(), 0) << "List of models is empty";
  count_ = dlr_models_.size();
  num_inputs_ = dlr_models_[0]->GetNumInputs();
  num_weights_ = dlr_models_[0]->GetNumWeights();
  num_outputs_ = dlr_models_.back()->GetNumOutputs();
  for (int i = 0; i < num_inputs_; i++) {
    input_names_.push_back(dlr_models_[0]->GetInputName(i));
    input_types_.push_back(dlr_models_[0]->GetInputType(i));
    input_shapes_.push_back(dlr_models_[0]->GetInputShape(i));
  }
  // Check previous model outputs and current model inputs compatibility
  for (int i = 1; i < count_; i++) {
    const DLRModelPtr prev_model = dlr_models_[i - 1];
    const DLRModelPtr curr_model = dlr_models_[i];
    CheckModelsCompatibility(prev_model, curr_model, i /*m1_id*/, false /*is_runtime_check*/);
  }
}

std::vector<std::string> PipelineModel::GetWeightNames() const {
  return dlr_models_[0]->GetWeightNames();
}

const char* PipelineModel::GetInputName(int index) const {
  return dlr_models_[0]->GetInputName(index);
}

const char* PipelineModel::GetInputType(int index) const {
  return dlr_models_[0]->GetInputType(index);
}

const int PipelineModel::GetInputDim(int index) const { return dlr_models_[0]->GetInputDim(index); }

const int64_t PipelineModel::GetInputSize(int index) const {
  return dlr_models_[0]->GetInputSize(index);
}

const char* PipelineModel::GetWeightName(int index) const {
  return dlr_models_[0]->GetWeightName(index);
}

void PipelineModel::SetInput(const char* name, const int64_t* shape, const void* input, int dim) {
  dlr_models_[0]->SetInput(name, shape, input, dim);
}

void PipelineModel::GetInput(const char* name, void* input) {
  dlr_models_[0]->GetInput(name, input);
}

void PipelineModel::GetOutputShape(int index, int64_t* shape) const {
  dlr_models_.back()->GetOutputShape(index, shape);
}

void PipelineModel::GetOutput(int index, void* out) { dlr_models_.back()->GetOutput(index, out); }

const void* PipelineModel::GetOutputPtr(int index) const {
  return dlr_models_.back()->GetOutputPtr(index);
}

void PipelineModel::GetOutputSizeDim(int index, int64_t* size, int* dim) {
  dlr_models_.back()->GetOutputSizeDim(index, size, dim);
}

const char* PipelineModel::GetOutputType(int index) const {
  return dlr_models_.back()->GetOutputType(index);
}

void PipelineModel::Run() {
  dlr_models_[0]->Run();
  for (int i = 1; i < count_; i++) {
    const DLRModelPtr prev_model = dlr_models_[i - 1];
    const DLRModelPtr curr_model = dlr_models_[i];
    CheckModelsCompatibility(prev_model, curr_model, i /*m1_id*/, true /*is_runtime_check*/);
    // for each model input
    for (int j = 0; j < curr_model->GetNumInputs(); j++) {
      const char* input_name = curr_model->GetInputName(j);
      // Get output shape of previous output.
      int64_t prev_output_size;
      int prev_output_dim;
      prev_model->GetOutputSizeDim(j, &prev_output_size, &prev_output_dim);
      std::vector<int64_t> prev_output_shape(prev_output_dim, -1);
      prev_model->GetOutputShape(j, prev_output_shape.data());
      const void* prev_model_output = prev_model->GetOutputPtr(j);
      curr_model->SetInput(input_name, prev_output_shape.data(), prev_model_output,
                           prev_output_dim);
    }
    curr_model->Run();
  }
}

void PipelineModel::SetNumThreads(int threads) {
  // Try to set Number of Threads to pipeline models
  // Ignore the errors in case some of the models do not support this feature.
  for (DLRModelPtr m : dlr_models_) {
    try {
      m->SetNumThreads(threads);
    } catch (dmlc::Error& e) {
      // ignore
    }
  }
}

void PipelineModel::UseCPUAffinity(bool use) {
  // Try to set UseCPUAffinity to pipeline models.
  // Ignore the errors in case some of the models do not support this feature.
  for (DLRModelPtr m : dlr_models_) {
    try {
      m->UseCPUAffinity(use);
    } catch (dmlc::Error& e) {
      // ignore
    }
  }
}

const char* PipelineModel::GetOutputName(const int index) const {
  return dlr_models_.back()->GetOutputName(index);
}

int PipelineModel::GetOutputIndex(const char* name) const {
  return dlr_models_.back()->GetOutputIndex(name);
}

void PipelineModel::GetOutputByName(const char* name, void* out) {
  return dlr_models_.back()->GetOutputByName(name, out);
}