/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   https://lammps.sandia.gov/, Sandia National Laboratories
   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract
   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
   certain rights in this software.  This software is distributed under
   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.
------------------------------------------------------------------------- */

/* ----------------------------------------------------------------------
   Contributing author: Pedro Antonio Santos Flórez (UNLV)
------------------------------------------------------------------------- */

#include "mliap_model_nn.h"
#include "pair_mliap.h"
#include "mliap_data.h"
#include "error.h"

#include "comm.h"
#include "memory.h"
#include "tokenizer.h"

#include <cmath>

using namespace LAMMPS_NS;

#define MAXLINE 1024

/* ---------------------------------------------------------------------- */

MLIAPModelNN::MLIAPModelNN(LAMMPS* lmp, char* coefffilename) :
  MLIAPModel(lmp, coefffilename)
{
  coeffelem = nullptr;
  nnodes = nullptr;
  activation = nullptr;
  scale = nullptr;
  if (coefffilename) read_coeffs(coefffilename);
}

/* ---------------------------------------------------------------------- */

MLIAPModelNN::~MLIAPModelNN()
{
    memory->destroy(coeffelem);
    memory->destroy(nnodes);
    memory->destroy(activation);
    memory->destroy(scale);
}

/* ----------------------------------------------------------------------
   get number of parameters
   ---------------------------------------------------------------------- */

int MLIAPModelNN::get_nparams()
{
  if (nparams == 0)
    if (ndescriptors == 0) error->all(FLERR,"ndescriptors not defined");

  return nparams;
}

void MLIAPModelNN::read_coeffs(char *coefffilename)
{

  // open coefficient file on proc 0

  FILE *fpcoeff;
  if (comm->me == 0) {
    fpcoeff = utils::open_potential(coefffilename,lmp,nullptr);
    if (fpcoeff == nullptr)
      error->one(FLERR,fmt::format("Cannot open MLIAPModel coeff file {}: {}",
                                   coefffilename,utils::getsyserror()));
  }

  char line[MAXLINE], *ptr, *tstr;
  int eof = 0;

  int n;
  int nwords = 0;
  while (nwords == 0) {
    if (comm->me == 0) {
      ptr = fgets(line,MAXLINE,fpcoeff);
      if (ptr == nullptr) {
        eof = 1;
        fclose(fpcoeff);
      } else n = strlen(line) + 1;
    }
    MPI_Bcast(&eof,1,MPI_INT,0,world);
    if (eof) break;
    MPI_Bcast(&n,1,MPI_INT,0,world);
    MPI_Bcast(line,n,MPI_CHAR,0,world);

    // strip comment, skip line if blank

    if ((ptr = strchr(line,'#'))) *ptr = '\0';
    nwords = utils::count_words(line);
  }
  if (nwords != 2)
    error->all(FLERR,"Incorrect format in MLIAPModel coefficient file");

  // words = ptrs to all words in line
  // strip single and double quotes from words

  try {
    ValueTokenizer coeffs(line);
    nelements = coeffs.next_int();
    nparams = coeffs.next_int();
  } catch (TokenizerException &e) {
    error->all(FLERR,fmt::format("Incorrect format in MLIAPModel coefficient "
                                 "file: {}",e.what()));
  }

  // set up coeff lists

  memory->create(coeffelem,nelements,nparams,"mliap_snap_model:coeffelem");

  int stats = 0;
  int ielem = 0;
  int l = 0;

  while (1) {
    if (comm->me == 0) {
      ptr = fgets(line,MAXLINE,fpcoeff);
      if (ptr == nullptr) {
        eof = 1;
        fclose(fpcoeff);
      } else n = strlen(line) + 1;
    }

    MPI_Bcast(&eof,1,MPI_INT,0,world);
    if (eof) break;
    MPI_Bcast(&n,1,MPI_INT,0,world);
    MPI_Bcast(line,n,MPI_CHAR,0,world);

    // strip comment, skip line if blank

    if ((ptr = strchr(line,'#'))) *ptr = '\0';
    nwords = utils::trim_and_count_words(line);
    if (nwords == 0) continue;

    if (stats == 0) { // Header NET
      tstr = strtok(line,"' \t\n\r\f");
      if (strncmp(tstr, "NET", 3) != 0) error->all(FLERR,"Incorrect format in MLIAPModel coefficient file");

      ndescriptors = atoi(strtok(nullptr,"' \t\n\r\f"));
      nlayers = atoi(strtok(nullptr,"' \t\n\r\f"));

      memory->create(activation,nlayers,"mliap_model:activation");
      memory->create(nnodes,nlayers,"mliap_model:nnodes");
      memory->create(scale,nelements,2,ndescriptors,"mliap_model:scale");

      for (int ilayer = 0; ilayer < nlayers; ilayer++) {
        tstr = strtok(NULL,"' \t\n\r\f");
        nnodes[ilayer] = atoi(strtok(NULL,"' \t\n\r\f"));

        if (strncmp(tstr, "linear", 6) == 0) activation[ilayer] = 0;
        else if (strncmp(tstr, "sigmoid", 7) == 0) activation[ilayer] = 1;
        else if (strncmp(tstr, "tanh", 4) == 0) activation[ilayer] = 2;
        else if (strncmp(tstr, "relu", 4) == 0) activation[ilayer] = 3;
        else activation[ilayer] = 4;
      }

      stats = 1;

    } else if (stats == 1) {
        scale[ielem][0][l] = atof(strtok(line,"' \t\n\r\f"));
        for (int icoeff = 1; icoeff < nwords; icoeff++) {
          scale[ielem][0][l+icoeff] = atof(strtok(nullptr,"' \t\n\r\f"));
        }
        l += nwords;
        if (l == ndescriptors) {
          stats = 2;
          l = 0;
        }

    } else if (stats == 2) {
        scale[ielem][1][l] = atof(strtok(line,"' \t\n\r\f"));
        for (int icoeff = 1; icoeff < nwords; icoeff++) {
          scale[ielem][1][l+icoeff] = atof(strtok(nullptr,"' \t\n\r\f"));
        }
        l += nwords;
        if (l == ndescriptors) {
          stats = 3;
          l = 0;
        }

    // set up coeff lists

    } else if (stats == 3) {
        if (nwords > 30) error->all(FLERR,"Incorrect format in MLIAPModel coefficient file");

        coeffelem[ielem][l] = atof(strtok(line,"' \t\n\r\f"));
        for (int icoeff = 1; icoeff < nwords; icoeff++) {
          coeffelem[ielem][l+icoeff] = atof(strtok(nullptr,"' \t\n\r\f"));
        }
        l += nwords;
        if (l == nparams) {
          stats = 1;
          l = 0;
          ielem++;
        }
    }
  }
}

/*  ----------------------------------------------------------------------
   Calculate model gradients w.r.t descriptors
   for each atom beta_i = dE(B_i)/dB_i
   ---------------------------------------------------------------------- */

void MLIAPModelNN::compute_gradients(MLIAPData* data)
{
  data->energy = 0.0;

  for (int ii = 0; ii < data->nlistatoms; ii++) {
      const int ielem = data->ielems[ii];
      const int nl = nlayers;

      double* coeffi = coeffelem[ielem];
      double** scalei = scale[ielem];
      double **nodes, **dnodes, **bnodes;

      nodes = new double*[nl];
      dnodes = new double*[nl];
      bnodes = new double*[nl];
      for (int l=0; l<nl; ++l) {
        nodes[l] = new double[nnodes[l]];
        dnodes[l] = new double[nnodes[l]];
        bnodes[l] = new double[nnodes[l]];
      }

      // forwardprop
      // input - hidden1

      for (int n=0; n < nnodes[0]; n++) {
        nodes[0][n] = 0;
        for (int icoeff = 0; icoeff < data->ndescriptors; icoeff++) {
          nodes[0][n] += coeffi[n*((data->ndescriptors)+1)+icoeff+1] *
            (data->descriptors[ii][icoeff] - scalei[0][icoeff]) /
            scalei[1][icoeff];
        }
        if (activation[0] == 1) {
          nodes[0][n] = sigm(nodes[0][n] +
                             coeffi[n*((data->ndescriptors)+1)],
                             dnodes[0][n]);
        }
        else if (activation[0] == 2) {
          nodes[0][n] = tanh(nodes[0][n] +
                             coeffi[n*((data->ndescriptors)+1)],
                             dnodes[0][n]);
        }
        else if (activation[0] == 3) {
          nodes[0][n] = relu(nodes[0][n] +
                             coeffi[n*((data->ndescriptors)+1)],
                             dnodes[0][n]);
        }
        else {
          nodes[0][n] += coeffi[n*((data->ndescriptors)+1)];
          dnodes[0][n] = 1;
        }
      }

      // hidden~output

      int k = 0;
      if (nl > 1) {
        k += ((data->ndescriptors)+1)*nnodes[0];
        for (int l=1; l < nl; l++) {
          for (int n=0; n < nnodes[l]; n++) {
            nodes[l][n] = 0;
            for (int j=0; j < nnodes[l-1]; j++) {
              nodes[l][n] += coeffi[k+n*(nnodes[l-1]+1)+j+1] *
                nodes[l-1][j];
            }
            if (activation[l] == 1) {
              nodes[l][n] = sigm(nodes[l][n] +
                                 coeffi[k+n*(nnodes[l-1]+1)],
                                 dnodes[l][n]);
            }
            else if (activation[l] == 2) {
              nodes[l][n] = tanh(nodes[l][n] +
                                 coeffi[k+n*(nnodes[l-1]+1)],
                                 dnodes[l][n]);
            }
            else if (activation[l] == 3) {
              nodes[l][n] = relu(nodes[l][n] +
                                 coeffi[k+n*(nnodes[l-1]+1)],
                                 dnodes[l][n]);
            }
            else {
              nodes[l][n] += coeffi[k+n*(nnodes[l-1]+1)];
              dnodes[l][n] = 1;
            }
          }
          k += (nnodes[l-1]+1)*nnodes[l];
        }
      }

      // backwardprop
      // output layer dnode initialized to 1.

      for (int n=0; n<nnodes[nl-1]; n++) {
        if (activation[nl-1] == 0) {
            bnodes[nl-1][n] = 1;
        }
        else {
            bnodes[nl-1][n] = dnodes[nl-1][n];
        }
      }

      if (nl > 1) {
        for (int l=nl-1; l>0; l--) {
          k -= (nnodes[l-1]+1)*nnodes[l];
          for (int n=0; n<nnodes[l-1]; n++) {
            bnodes[l-1][n] = 0;
            for (int j=0; j<nnodes[l]; j++) {
              bnodes[l-1][n] += coeffi[k+j*(nnodes[l-1]+1)+n+1] *
                bnodes[l][j];
            }
            if (activation[l-1] >= 1) {
              bnodes[l-1][n] *= dnodes[l-1][n];
            }
          }
        }
      }

      for (int icoeff = 0; icoeff < data->ndescriptors; icoeff++) {
        data->betas[ii][icoeff] = 0;
        for (int j=0; j<nnodes[0]; j++) {
          data->betas[ii][icoeff] +=
            coeffi[j*((data->ndescriptors)+1)+icoeff+1] *
            bnodes[0][j];
        }
        data->betas[ii][icoeff] = data->betas[ii][icoeff]/scalei[1][icoeff];
      }

      if (data->eflag) {

      // energy of atom I (E_i)

        double etmp = nodes[nl-1][0];

        data->energy += etmp;
        data->eatoms[ii] = etmp;
      }
      // Deleting the variables

      for (int n=0; n<nl; n++) {
        delete nodes[n];
        delete dnodes[n];
        delete bnodes[n];
      }

      delete[] nodes;
      delete[] dnodes;
      delete[] bnodes;

  }
}

/* ----------------------------------------------------------------------
   Calculate model double gradients w.r.t descriptors and parameters
   for each atom energy gamma_lk = d2E(B)/dB_k/dsigma_l,
   where sigma_l is a parameter, B_k a descriptor,
   and atom subscript i is omitted

   gamma is in CSR format:
      nnz = number of non-zero values
      gamma_row_index[inz] = l indices, 0 <= l < nparams
      gamma_col_indexiinz] = k indices, 0 <= k < ndescriptors
      gamma[i][inz] = non-zero values, 0 <= inz < nnz

   egradient is derivative of energy w.r.t. parameters
   ---------------------------------------------------------------------- */

void MLIAPModelNN::compute_gradgrads(class MLIAPData * /*data*/)
{
  error->all(FLERR,"compute_gradgrads not implemented");
}

/* ----------------------------------------------------------------------
   calculate gradients of forces w.r.t. parameters
   egradient is derivative of energy w.r.t. parameters
   ---------------------------------------------------------------------- */

void MLIAPModelNN::compute_force_gradients(class MLIAPData * /*data*/)
{
  error->all(FLERR,"compute_force_gradients not implemented");
}

/* ----------------------------------------------------------------------
   count the number of non-zero entries in gamma matrix
   ---------------------------------------------------------------------- */

int MLIAPModelNN::get_gamma_nnz(class MLIAPData * /*data*/)
{
  // todo: get_gamma_nnz
  return 0;
}

double MLIAPModelNN::memory_usage()
{
  double bytes = 0;

  bytes += (double)nelements*nparams*sizeof(double);         // coeffelem
  bytes += (double)nelements*2*ndescriptors*sizeof(double);  // scale
  bytes += (int)nlayers*sizeof(int);                         // nnodes
  bytes += (int)nlayers*sizeof(int);                         // activation
  return bytes;
}