// 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 "network.hpp"
#include <assert.h>

Network::Network() {
  nlayers_global = 0;
  nlayers_local = 0;
  nchannels = 0;
  dt = 0.0;
  loss = 0.0;
  accuracy = 0.0;

  startlayerID = 0;
  endlayerID = 0;

  ndesign_local = 0;
  ndesign_global = 0;
  ndesign_layermax = 0;

  design = NULL;
  gradient = NULL;

  layers = NULL;
  openlayer = NULL;
  layer_left = NULL;
  layer_right = NULL;

  comm = MPI_COMM_WORLD;
}

void Network::createNetworkBlock(int StartLayerID, int EndLayerID,
                                 Config *config, MPI_Comm Comm) {
  /* Initilizize */
  startlayerID = StartLayerID;
  endlayerID = EndLayerID;
  nlayers_local = endlayerID - startlayerID + 1;
  nlayers_global = config->nlayers;
  nchannels = config->nchannels;
  dt = (config->T) / (MyReal)(config->nlayers - 2);  // nlayers-2 = nhiddenlayers
  comm = Comm;

  /* --- Create the layers --- */
  ndesign_local = 0;

  if (startlayerID == 0)  // Opening layer
  {
    /* Create the opening layer */
    int index = -1;
    openlayer = createLayer(index, config);
    ndesign_local += openlayer->getnDesign();
    // printf("Create opening layer %d, ndesign_local %d \n", index,
    // openlayer->getnDesign());
  }

  layers = new Layer *[nlayers_local];  // Intermediate and classification layer
  for (int ilayer = startlayerID; ilayer <= endlayerID; ilayer++) {
    /* Create a layer at time step ilayer. Local storage at ilayer -
     * startlayerID */
    int storeID = getLocalID(ilayer);
    layers[storeID] = createLayer(ilayer, config);
    ndesign_local += layers[storeID]->getnDesign();
    // printf("creating hidden/class layer %d/%d, ndesign_local%d\n", ilayer,
    // nlayers_local, layers[storeID]->getnDesign());
  }

  /* Create left neighbouring layer */
  int leftID = startlayerID - 1;
  layer_left = createLayer(leftID, config);

  /* Create right neighbrouing layer */
  int rightID = endlayerID + 1;
  layer_right = createLayer(rightID, config);

  /* Sum up global number of design vars */
  MPI_Allreduce(&ndesign_local, &ndesign_global, 1, MPI_INT, MPI_SUM, comm);

  /* Store maximum number of designs over all layers layermax */
  ndesign_layermax = computeLayermax();

  /* Allocate memory for network design and gradient variables */
  design = new MyReal[ndesign_local];
  gradient = new MyReal[ndesign_local];

  /* Set the memory locations for all layers */
  int istart = 0;
  if (openlayer != NULL)  // Openlayer on first processor
  {
    openlayer->setMemory(&(design[istart]), &(gradient[istart]));
    istart += openlayer->getnDesign();
  }
  for (int ilayer = startlayerID; ilayer <= endlayerID;
       ilayer++)  // intermediate and hidden layers
  {
    layers[getLocalID(ilayer)]->setMemory(&(design[istart]),
                                          &(gradient[istart]));
    istart += layers[getLocalID(ilayer)]->getnDesign();
  }

  /* left anr right neighbouring layer design, if exists */
  if (layer_left != NULL) {
    MyReal *left_design = new MyReal[layer_left->getnDesign()];
    MyReal *left_gradient = new MyReal[layer_left->getnDesign()];
    layer_left->setMemory(left_design, left_gradient);
  }
  /* Create and initialize right neighbouring layer design, if exists */
  if (layer_right != NULL) {
    MyReal *right_design = new MyReal[layer_right->getnDesign()];
    MyReal *right_gradient = new MyReal[layer_right->getnDesign()];
    layer_right->setMemory(right_design, right_gradient);
  }
}

Network::~Network() {
  /* Delete openlayer */
  if (openlayer != NULL) delete openlayer;

  /* Delete intermediate and classification layers */
  for (int ilayer = 0; ilayer < nlayers_local; ilayer++) {
    delete layers[ilayer];
  }
  delete[] layers;

  /* Delete design and gradient */
  delete[] design;
  delete[] gradient;

  /* Delete neighbouring layer information */
  if (layer_left != NULL) {
    delete[] layer_left->getWeights();
    delete[] layer_left->getWeightsBar();
    delete layer_left;
  }

  if (layer_right != NULL) {
    delete[] layer_right->getWeights();
    delete[] layer_right->getWeightsBar();
    delete layer_right;
  }
}

int Network::getnChannels() { return nchannels; }

int Network::getnLayersGlobal() { return nlayers_global; }

MyReal Network::getDT() { return dt; }

int Network::getLocalID(int ilayer) {
  int idx = ilayer - startlayerID;
  return idx;
}

MyReal Network::getLoss() { return loss; }

MyReal Network::getAccuracy() { return accuracy; }

int Network::getnDesignLocal() { return ndesign_local; }

int Network::getnDesignGlobal() { return ndesign_global; }

MyReal *Network::getDesign() { return design; }

MyReal *Network::getGradient() { return gradient; }

int Network::getStartLayerID() { return startlayerID; }
int Network::getEndLayerID() { return endlayerID; }

MPI_Comm Network::getComm() { return comm; }

Layer *Network::createLayer(int index, Config *config) {
  Layer *layer = 0;
  if (index == -1)  // Opening layer
  {
    switch (config->network_type) {
      case DENSE:
        if (config->weights_open_init == 0.0) {
          layer = new OpenExpandZero(config->nfeatures, nchannels);
        } else {
          layer = new OpenDenseLayer(config->nfeatures, nchannels,
                                     config->activation, config->gamma_tik);
        }
        break;
      case CONVOLUTIONAL:
        /**< (Weight_open_init == 0.0) not needed for convolutional layers*/
        if (config->openlayer_type == 0) {
          layer = new OpenConvLayer(config->nfeatures, nchannels);
        } else if (config->openlayer_type == 1) {
          layer = new OpenConvLayerMNIST(config->nfeatures, nchannels);
        }
        break;
    }
  } else if (0 <= index && index < nlayers_global - 2)  // Intermediate layer
  {
    switch (config->network_type) {
      case DENSE:
        layer =
            new DenseLayer(index, nchannels, nchannels, dt, config->activation,
                           config->gamma_tik, config->gamma_ddt);
        break;
      case CONVOLUTIONAL:
        // TODO: Fix
        int convolution_size = 3;
        layer =
            new ConvLayer(index, nchannels, nchannels, convolution_size,
                          nchannels / config->nfeatures, dt, config->activation,
                          config->gamma_tik, config->gamma_ddt);
        break;
    }
  } else if (index == nlayers_global - 2)  // Classification layer
  {
    layer = new ClassificationLayer(index, nchannels, config->nclasses,
                                    config->gamma_class);
  } else {
    layer = NULL;
  }

  return layer;
}

Layer *Network::getLayer(int layerindex) {
  Layer *layer;

  if (layerindex == -1)  // opening layer
  {
    layer = openlayer;
  } else if (layerindex == startlayerID - 1) {
    layer = layer_left;
  } else if (startlayerID <= layerindex && layerindex <= endlayerID) {
    layer = layers[getLocalID(layerindex)];
  } else if (layerindex == endlayerID + 1) {
    layer = layer_right;
  } else {
    layer = NULL;
  }

  return layer;
}

int Network::getnDesignLayermax() { return ndesign_layermax; }

int Network::computeLayermax() {
  int ndesignlayer;
  int max = 0;

  /* Loop over all local layers */
  for (int ilayer = startlayerID; ilayer <= endlayerID; ilayer++) {
    if (ilayer < nlayers_global - 2)  // excludes classification layer
    {
      /* Update maximum */
      ndesignlayer = layers[getLocalID(ilayer)]->getnDesign();
      if (ndesignlayer > max) max = ndesignlayer;
    }
  }

  /* Get maximum over all local layer blocks */
  int mymax = max;
  MPI_Allreduce(&mymax, &max, 1, MPI_INT, MPI_MAX, comm);

  return max;
}

void Network::setInitialDesign(Config *config) {
  MyReal factor;
  MyReal *design_init;
  char filename[255];
  int myid;
  MPI_Comm_rank(comm, &myid);

  /* Initialize design with random numbers (do on one processor and scatter for
   * scaling test) */
  if (myid == 0) {
    srand(1.0);
    design_init = new MyReal[ndesign_global];
    for (int i = 0; i < ndesign_global; i++) {
      design_init[i] = (MyReal)rand() / ((MyReal)RAND_MAX);
    }
  }
  /* Scatter initial design to all processors */
  MPI_ScatterVector(design_init, design, ndesign_local, 0, comm);

  /* Scale the initial design by a factor and read from file, if set */

  /* Opening layer on first processor */
  if (startlayerID == 0) {
    /* Scale design by the factor */
    factor = config->weights_open_init;
    openlayer->scaleDesign(factor);

    /* if set, overwrite opening design from file */
    if (strcmp(config->weightsopenfile, "NONE") != 0) {
      sprintf(filename, "%s/%s", config->datafolder, config->weightsopenfile);
      read_vector(filename, openlayer->getWeights(), openlayer->getnDesign());
    }
  }

  /* Intermediate (hidden) and classification layers */
  for (int ilayer = startlayerID; ilayer <= endlayerID; ilayer++) {
    if (ilayer < nlayers_global - 1)  // Intermediate layer
    {
      factor = config->weights_init;
    } else  // Classification layer
    {
      factor = config->weights_class_init;
    }

    /* Set memory location and scale the current design by the factor */
    int storeID = getLocalID(ilayer);
    layers[storeID]->scaleDesign(factor);

    /* if set, overwrite classification design from file */
    if (ilayer == nlayers_global - 1) {
      if (strcmp(config->weightsclassificationfile, "NONE") != 0) {
        sprintf(filename, "%s/%s", config->datafolder,
                config->weightsclassificationfile);
        read_vector(filename, layers[storeID]->getWeights(),
                    layers[storeID]->getnDesign());
      }
    }
  }

  /* Communicate the neighbours across processors */
  MPI_CommunicateNeighbours(comm);

  if (myid == 0) delete[] design_init;
}

void Network::MPI_CommunicateNeighbours(MPI_Comm comm) {
  int myid, comm_size;
  MPI_Comm_rank(comm, &myid);
  MPI_Comm_size(comm, &comm_size);
  MPI_Request sendlastreq, recvlastreq;
  MPI_Request sendfirstreq, recvfirstreq;
  MPI_Status status;

  /* Allocate buffers */
  int size_left = -1;
  int size_right = -1;

  MyReal *sendlast = 0;
  MyReal *recvlast = 0;
  MyReal *sendfirst = 0;
  MyReal *recvfirst = 0;

  /* --- All but the first process receive the last layer from left neighbour
   * --- */
  if (myid > 0) {
    /* Receive from left neighbour */
    int source = myid - 1;

    size_left = layer_left->getnDesign();
    recvlast = new MyReal[size_left];

    MPI_Irecv(recvlast, size_left, MPI_MyReal, source, 0, comm, &recvlastreq);
  }

  /* --- All but the last process sent their last layer to right neighbour ---
   */
  if (myid < comm_size - 1) {
    size_left = layers[getLocalID(endlayerID)]->getnDesign();
    sendlast = new MyReal[size_left];

    /* Pack the last layer into a buffer */
    layers[getLocalID(endlayerID)]->packDesign(sendlast, size_left);

    /* Send to right neighbour */
    int receiver = myid + 1;
    MPI_Isend(sendlast, size_left, MPI_MyReal, receiver, 0, comm, &sendlastreq);
  }

  /* --- All but the last processor recv the first layer from the right
   * neighbour --- */
  if (myid < comm_size - 1) {
    /* Receive from right neighbour */
    int source = myid + 1;

    size_right = layer_right->getnDesign();
    recvfirst = new MyReal[size_right];

    MPI_Irecv(recvfirst, size_right, MPI_MyReal, source, 1, comm,
              &recvfirstreq);
  }

  /* --- All but the first processor send their first layer to the left
   * neighbour --- */
  if (myid > 0) {
    size_right = layers[getLocalID(startlayerID)]->getnDesign();
    sendfirst = new MyReal[size_right];

    /* Pack the first layer into a buffer */
    layers[getLocalID(startlayerID)]->packDesign(sendfirst, size_right);

    /* Send to left neighbour */
    int receiver = myid - 1;
    MPI_Isend(sendfirst, size_right, MPI_MyReal, receiver, 1, comm,
              &sendfirstreq);
  }

  /* Wait to finish up communication */
  if (myid > 0) MPI_Wait(&recvlastreq, &status);
  if (myid < comm_size - 1) MPI_Wait(&sendlastreq, &status);
  if (myid < comm_size - 1) MPI_Wait(&recvfirstreq, &status);
  if (myid > 0) MPI_Wait(&sendfirstreq, &status);

  /* Unpack and store the left received layer */
  if (myid > 0) {
    layer_left->unpackDesign(recvlast);
  }

  /* Unpack and store the right received layer */
  if (myid < comm_size - 1) {
    layer_right->unpackDesign(recvfirst);
  }

  /* Free the buffer */
  if (sendlast != 0) delete[] sendlast;
  if (recvlast != 0) delete[] recvlast;
  if (sendfirst != 0) delete[] sendfirst;
  if (recvfirst != 0) delete[] recvfirst;
}

void Network::evalClassification(DataSet *data, MyReal **state, int output) {
  MyReal *tmpstate = new MyReal[nchannels];

  int class_id;
  int success, success_local;
  FILE *classfile;
  ClassificationLayer *classificationlayer;

  /* Get classification layer */
  classificationlayer =
      dynamic_cast<ClassificationLayer *>(getLayer(nlayers_global - 2));
  if (classificationlayer == NULL) {
    printf("\n ERROR: Network can't access classification layer!\n\n");
    exit(1);
  }

  /* open file for printing predicted file */
  if (output) classfile = fopen("classprediction.dat", "w");

  loss = 0.0;
  accuracy = 0.0;
  success = 0;
  for (int iex = 0; iex < data->getnBatch(); iex++) {
    /* Copy values so that they are not overwrittn (they are needed for
     * adjoint)*/
    for (int ic = 0; ic < nchannels; ic++) {
      tmpstate[ic] = state[iex][ic];
    }
    /* Apply classification on tmpstate */
    classificationlayer->setLabel(data->getLabel(iex));
    classificationlayer->applyFWD(tmpstate);
    /* Evaluate Loss */
    loss += classificationlayer->crossEntropy(tmpstate);
    success_local = classificationlayer->prediction(tmpstate, &class_id);
    success += success_local;
    if (output) fprintf(classfile, "%d   %d\n", class_id, success_local);
  }
  loss = 1. / data->getnBatch() * loss;
  accuracy = 100.0 * ((MyReal)success) / data->getnBatch();
  // printf("Classification %d: %1.14e using layer %1.14e state %1.14e
  // tmpstate[0] %1.14e\n", getIndex(), loss, weights[0], state[1][1],
  // tmpstate[0]);

  if (output) fclose(classfile);
  if (output) printf("Prediction file written: classprediction.dat\n");

  delete[] tmpstate;
}

void Network::evalClassification_diff(DataSet *data, MyReal **primalstate,
                                      MyReal **adjointstate,
                                      int compute_gradient) {
  MyReal *tmpstate = new MyReal[nchannels];
  ClassificationLayer *classificationlayer;

  /* Get classification layer */
  classificationlayer =
      dynamic_cast<ClassificationLayer *>(getLayer(nlayers_global - 2));
  if (classificationlayer == NULL) {
    printf("\n ERROR: Network can't access classification layer!\n\n");
    exit(1);
  }

  int nbatch = data->getnBatch();
  MyReal loss_bar = 1. / nbatch;

  for (int iex = 0; iex < nbatch; iex++) {
    /* Recompute the Classification */
    for (int ic = 0; ic < nchannels; ic++) {
      tmpstate[ic] = primalstate[iex][ic];
    }
    classificationlayer->setLabel(data->getLabel(iex));
    classificationlayer->applyFWD(tmpstate);

    /* Derivative of Loss and classification. */
    classificationlayer->crossEntropy_diff(tmpstate, adjointstate[iex],
                                           loss_bar);
    classificationlayer->applyBWD(primalstate[iex], adjointstate[iex],
                                  compute_gradient);
  }
  // printf("Classification_diff %d using layer %1.14e state %1.14e tmpstate
  // %1.14e biasbar[dimOut-1] %1.14e\n", getIndex(), weights[0],
  // primalstate[1][1], tmpstate[0], bias_bar[dim_Out-1]);

  delete[] tmpstate;
}

void Network::updateDesign(MyReal stepsize, MyReal *direction, MPI_Comm comm) {
  /* Update design locally on this network-block */
  for (int id = 0; id < ndesign_local; id++) {
    design[id] += stepsize * direction[id];
  }

  /* Communicate design across neighbouring processors (ghostlayers) */
  MPI_CommunicateNeighbours(comm);
}