// Copyright
//
// 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.
//
// Underlying paper:
//
// Layer-Parallel Training of Deep Residual Neural Networks
// S. Guenther, L. Ruthotto, J.B. Schroder, E.C. Czr, and N.R. Gauger
//
// Download: https://arxiv.org/pdf/1812.04352.pdf
//
#include "dataset.hpp"

DataSet::DataSet() {
  nelements = 0;
  nfeatures = 0;
  nlabels = 0;
  nbatch = 0;
  MPIsize = 0;
  MPIrank = 0;
  navail = 0;

  examples = NULL;
  labels = NULL;
  batchIDs = NULL;
  availIDs = NULL;
}

void DataSet::initialize(int nElements, int nFeatures, int nLabels, int nBatch,
                         MPI_Comm comm) {
  nelements = nElements;
  nfeatures = nFeatures;
  nlabels = nLabels;
  nbatch = nBatch;
  navail = nelements;

  MPI_Comm_rank(comm, &MPIrank);
  MPI_Comm_size(comm, &MPIsize);

  /* Sanity check */
  if (nbatch > nelements) nbatch = nelements;

  /* Allocate feature vectors on first processor */
  if (MPIrank == 0) {
    examples = new MyReal *[nelements];
    for (int ielem = 0; ielem < nelements; ielem++) {
      examples[ielem] = new MyReal[nfeatures];
    }
  }
  /* Allocate label vectors on last processor */
  if (MPIrank == MPIsize - 1) {
    labels = new MyReal *[nelements];
    for (int ielem = 0; ielem < nelements; ielem++) {
      labels[ielem] = new MyReal[nlabels];
    }
  }

  /* Allocate and initialize availIDs and batchIDs on first and last processor
   */
  if (MPIrank == 0 || MPIrank == MPIsize - 1) {
    availIDs = new int[nelements];  // all elements
    batchIDs = new int[nbatch];

    /* Initialize available ID with identity */
    for (int idx = 0; idx < nelements; idx++) {
      availIDs[idx] = idx;
    }

    /* Initialize the batch with identity */
    for (int idx = 0; idx < nbatch; idx++) {
      batchIDs[idx] = idx;
    }
  }
}

DataSet::~DataSet() {
  /* Deallocate feature vectors on first processor */
  if (examples != NULL) {
    for (int ielem = 0; ielem < nelements; ielem++) {
      delete[] examples[ielem];
    }
    delete[] examples;
  }

  /* Deallocate label vectors on last processor */
  if (labels != NULL) {
    for (int ielem = 0; ielem < nelements; ielem++) {
      delete[] labels[ielem];
    }
    delete[] labels;
  }

  if (availIDs != NULL) delete[] availIDs;
  if (batchIDs != NULL) delete[] batchIDs;
}

int DataSet::getnBatch() { return nbatch; }

MyReal *DataSet::getExample(int id) {
  if (examples == NULL) return NULL;

  return examples[batchIDs[id]];
}

MyReal *DataSet::getLabel(int id) {
  if (labels == NULL) return NULL;

  return labels[batchIDs[id]];
}

void DataSet::readData(const char *datafolder, const char *examplefile,
                       const char *labelfile) {
  char examplefilename[255], labelfilename[255];

  /* Set the file names  */
  sprintf(examplefilename, "%s/%s", datafolder, examplefile);
  sprintf(labelfilename, "%s/%s", datafolder, labelfile);

  /* Read feature vectors on first processor */
  if (MPIrank == 0)
    read_matrix(examplefilename, examples, nelements, nfeatures);

  /* Read label vectors on last processor) */
  if (MPIrank == MPIsize - 1)
    read_matrix(labelfilename, labels, nelements, nlabels);
}

void DataSet::selectBatch(int batch_type, MPI_Comm comm) {
  int irand, rand_range;
  int tmp;
  MPI_Request sendreq, recvreq;
  MPI_Status status;

  switch (batch_type) {
    case DETERMINISTIC:
      /* Do nothing, keep the batch fixed. */
      break;

    case STOCHASTIC:

      /* Randomly choose a batch on first processor, send to last processor */
      if (MPIrank == 0) {
        /* Fill the batchID vector with randomly generated integer */
        rand_range = navail - 1;
        for (int ibatch = 0; ibatch < nbatch; ibatch++) {
          /* Generate a new random index in [0,range] */
          irand = (int)((((double)rand()) / (double)RAND_MAX) * rand_range);

          /* Set the batchID */
          batchIDs[ibatch] = availIDs[irand];

          /* Remove the ID from available IDs (by swapping it with the last
           * available id and reducing the range) */
          tmp = availIDs[irand];
          availIDs[irand] = availIDs[rand_range];
          availIDs[rand_range] = tmp;
          rand_range--;
        }

        /* Send to the last processor */
        int receiver = MPIsize - 1;
        MPI_Isend(batchIDs, nbatch, MPI_INT, receiver, 0, comm, &sendreq);
      }

      /* Receive the batch IDs on last processor */
      if (MPIrank == MPIsize - 1) {
        int source = 0;
        MPI_Irecv(batchIDs, nbatch, MPI_INT, source, 0, comm, &recvreq);
      }

      /* Wait to finish communication */
      if (MPIrank == 0) MPI_Wait(&sendreq, &status);
      if (MPIrank == MPIsize - 1) MPI_Wait(&recvreq, &status);

      break;  // break switch statement
  }
}

void DataSet::printBatch() {
  if (batchIDs != NULL)  // only first and last processor
  {
    printf("%d:\n", MPIrank);
    for (int ibatch = 0; ibatch < nbatch; ibatch++) {
      printf("%d, %04d\n", ibatch, batchIDs[ibatch]);
    }
  }
}