/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   https://www.lammps.org/, Sandia National Laboratories
   LAMMPS development team: developers@lammps.org

   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: Aidan Thompson (SNL)
------------------------------------------------------------------------- */

#include "mliap_data.h"

#include "atom.h"
#include "error.h"
#include "memory.h"
#include "mliap_descriptor.h"
#include "mliap_model.h"
#include "neigh_list.h"

using namespace LAMMPS_NS;

MLIAPData::MLIAPData(LAMMPS *lmp, int gradgradflag_in, int *map_in, class MLIAPModel *model_in,
                     class MLIAPDescriptor *descriptor_in, class PairMLIAP *pairmliap_in) :
    Pointers(lmp),
    f(nullptr), gradforce(nullptr), betas(nullptr), descriptors(nullptr), eatoms(nullptr),
    gamma(nullptr), gamma_row_index(nullptr), gamma_col_index(nullptr), egradient(nullptr),
    numneighs(nullptr), iatoms(nullptr), ielems(nullptr), itypes(nullptr), pair_i(nullptr),
    jatoms(nullptr), jelems(nullptr), elems(nullptr), lmp_firstneigh(nullptr), rij(nullptr),
    graddesc(nullptr), model(nullptr), descriptor(nullptr), list(nullptr)
{
  gradgradflag = gradgradflag_in;
  map = map_in;
  model = model_in;
  descriptor = descriptor_in;
  pairmliap = pairmliap_in;

  ndescriptors = descriptor->ndescriptors;
  nelements = descriptor->nelements;
  nparams = model->get_nparams();

  gamma_nnz = model->get_gamma_nnz(this);
  ndims_force = 3;
  ndims_virial = 6;
  yoffset = nparams * nelements;
  zoffset = 2 * yoffset;
  natoms = atom->natoms;

  // must check before assigning bigint expression to regular int

  if (1 + ndims_force * natoms + ndims_virial > MAXSMALLINT)
    error->all(FLERR, "Too many atoms for MLIAP package");

  size_array_rows = 1 + ndims_force * natoms + ndims_virial;
  size_array_cols = nparams * nelements + 1;
  size_gradforce = ndims_force * nparams * nelements;

  nlistatoms_max = 0;
  natomneigh_max = 0;
  nneigh_max = 0;
  nmax = 0;
  natomgamma_max = 0;
}

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

MLIAPData::~MLIAPData()
{
  memory->destroy(betas);
  memory->destroy(descriptors);
  memory->destroy(eatoms);
  memory->destroy(gamma_row_index);
  memory->destroy(gamma_col_index);
  memory->destroy(gamma);
  memory->destroy(egradient);
  memory->destroy(gradforce);

  memory->destroy(iatoms);
  memory->destroy(pair_i);
  memory->destroy(ielems);
  memory->destroy(itypes);
  memory->destroy(numneighs);
  memory->destroy(jatoms);
  memory->destroy(lmp_firstneigh);
  memory->destroy(jelems);
  memory->destroy(elems);
  memory->destroy(rij);
  memory->destroy(graddesc);
}

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

void MLIAPData::init()
{
  memory->create(egradient, nelements * nparams, "MLIAPData:egradient");
}

/* ----------------------------------------------------------------------
   generate neighbor arrays
------------------------------------------------------------------------- */

void MLIAPData::generate_neighdata(NeighList *list_in, int eflag_in, int vflag_in)
{
  list = list_in;
  f = atom->f;
  double **x = atom->x;
  int *type = atom->type;

  int *ilist = list->ilist;
  int *numneigh = list->numneigh;
  int **firstneigh = list->firstneigh;

  int nall = atom->nlocal + atom->nghost;
  nlocal = atom->nlocal;
  ntotal = nall;

  // grow nmax gradforce, elems arrays if necessary

  if (atom->nmax > nmax) {
    nmax = atom->nmax;
    memory->grow(elems, nmax, "MLIAPData:elems");
    if (gradgradflag > -1) memory->grow(gradforce, nmax, size_gradforce, "MLIAPData:gradforce");
  }

  // clear gradforce, elems arrays

  for (int i = 0; i < nall; i++) {
    elems[i] = 0;
    if (gradgradflag > -1) {
      for (int j = 0; j < size_gradforce; j++) { gradforce[i][j] = 0.0; }
    }
  }

  // grow arrays if necessary

  nlistatoms = list->inum;
  if (nlistatoms_max < nlistatoms) {
    memory->grow(betas, nlistatoms, ndescriptors, "MLIAPData:betas");
    memory->grow(descriptors, nlistatoms, ndescriptors, "MLIAPData:descriptors");
    memory->grow(eatoms, nlistatoms, "MLIAPData:eatoms");
    nlistatoms_max = nlistatoms;
  }

  // grow gamma arrays if necessary

  if (gradgradflag == 1) {
    if (natomgamma_max < nlistatoms) {
      memory->grow(gamma_row_index, nlistatoms, gamma_nnz, "MLIAPData:gamma_row_index");
      memory->grow(gamma_col_index, nlistatoms, gamma_nnz, "MLIAPData:gamma_col_index");
      memory->grow(gamma, nlistatoms, gamma_nnz, "MLIAPData:gamma");
      natomgamma_max = nlistatoms;
    }
  }

  grow_neigharrays();

  npairs = 0;
  int ij = 0;
  for (int ii = 0; ii < natomneigh; ii++) {
    const int i = ilist[ii];

    const double xtmp = x[i][0];
    const double ytmp = x[i][1];
    const double ztmp = x[i][2];
    const int itype = type[i];
    const int ielem = map[itype];

    int *jlist = firstneigh[i];
    const int jnum = numneigh[i];

    int ninside = 0;
    for (int jj = 0; jj < jnum; jj++) {
      int j = jlist[jj];
      j &= NEIGHMASK;
      const double delx = x[j][0] - xtmp;
      const double dely = x[j][1] - ytmp;
      const double delz = x[j][2] - ztmp;
      const double rsq = delx * delx + dely * dely + delz * delz;
      int jtype = type[j];
      const int jelem = map[jtype];

      if (rsq < descriptor->cutsq[ielem][jelem]) {
        pair_i[ij] = i;
        jatoms[ij] = j;
        jelems[ij] = jelem;
        rij[ij][0] = delx;
        rij[ij][1] = dely;
        rij[ij][2] = delz;
        lmp_firstneigh[ii][ninside] = firstneigh[i][jj];
        ij++;
        ninside++;
      }
    }
    iatoms[ii] = i;
    ielems[ii] = ielem;
    itypes[ii] = itype;
    numneighs[ii] = ninside;
    npairs += ninside;
  }

  for (int i = 0; i < nall; i++) {
    const int itype = type[i];
    elems[i] = map[itype];
  }

  eflag = eflag_in;
  vflag = vflag_in;
}

/* ----------------------------------------------------------------------
   grow neighbor arrays to handle all neighbors
------------------------------------------------------------------------- */

void MLIAPData::grow_neigharrays()
{

  // grow neighbor atom arrays if necessary
  natomneigh = list->inum;
  if (list->ghost == 1) natomneigh += list->gnum;
  if (natomneigh_max < natomneigh) {
    memory->grow(iatoms, natomneigh, "MLIAPData:iatoms");
    memory->grow(ielems, natomneigh, "MLIAPData:ielems");
    memory->grow(itypes, natomneigh, "MLIAPData:itypes");
    memory->grow(numneighs, natomneigh, "MLIAPData:numneighs");
    memory->grow(lmp_firstneigh, natomneigh, nneigh_max, "MLIAPData:lmp_firstneigh");
    natomneigh_max = natomneigh;
  }

  // grow neighbor arrays if necessary

  int *ilist = list->ilist;
  int *numneigh = list->numneigh;
  int **firstneigh = list->firstneigh;
  double **x = atom->x;
  int *type = atom->type;

  int nneigh = 0;
  for (int ii = 0; ii < natomneigh; ii++) {
    const int i = ilist[ii];

    const double xtmp = x[i][0];
    const double ytmp = x[i][1];
    const double ztmp = x[i][2];
    const int itype = type[i];
    const int ielem = map[itype];

    int *jlist = firstneigh[i];
    const int jnum = numneigh[i];

    int ninside = 0;
    for (int jj = 0; jj < jnum; jj++) {
      int j = jlist[jj];
      j &= NEIGHMASK;
      const double delx = x[j][0] - xtmp;
      const double dely = x[j][1] - ytmp;
      const double delz = x[j][2] - ztmp;
      const double rsq = delx * delx + dely * dely + delz * delz;
      int jtype = type[j];
      const int jelem = map[jtype];
      if (rsq < descriptor->cutsq[ielem][jelem]) ninside++;
    }
    nneigh += ninside;
  }

  if (nneigh_max < nneigh) {
    memory->grow(pair_i, nneigh, "MLIAPData:pair_i");
    memory->grow(jatoms, nneigh, "MLIAPData:jatoms");
    memory->grow(lmp_firstneigh, natomneigh, nneigh, "MLIAPData:lmp_firstneigh");
    memory->grow(jelems, nneigh, "MLIAPData:jelems");
    memory->grow(rij, nneigh, 3, "MLIAPData:rij");
    if (gradgradflag == 0) memory->grow(graddesc, nneigh, ndescriptors, 3, "MLIAPData:graddesc");
    nneigh_max = nneigh;
  }
}

double MLIAPData::memory_usage()
{
  double bytes = 0.0;

  bytes += (double) nelements * nparams * sizeof(double);      // egradient
  bytes += (double) nmax * size_gradforce * sizeof(double);    // gradforce
  bytes += (double) nmax * sizeof(int);                        // elems

  if (gradgradflag == 1) {
    bytes += (double) natomgamma_max * gamma_nnz * sizeof(int);       //gamma_row_index
    bytes += (double) natomgamma_max * gamma_nnz * sizeof(int);       // gamma_col_index
    bytes += (double) natomgamma_max * gamma_nnz * sizeof(double);    // gamma
  }

  bytes += (double) nlistatoms * ndescriptors * sizeof(int);    // betas
  bytes += (double) nlistatoms * ndescriptors * sizeof(int);    // descriptors
  bytes += (double) nlistatoms * sizeof(double);                // eatoms

  bytes += (double) natomneigh_max * sizeof(int);    // iatoms
  bytes += (double) natomneigh_max * sizeof(int);    // ielems
  bytes += (double) natomneigh_max * sizeof(int);    // itypes
  bytes += (double) natomneigh_max * sizeof(int);    // numneighs

  bytes += (double) nneigh_max * sizeof(int);                     // pair_i
  bytes += (double) nneigh_max * sizeof(int);                     // jatoms
  bytes += (double) nneigh_max * sizeof(int);                     // jelems
  bytes += (double) nneigh_max * natomneigh_max * sizeof(int);    // lmp_firstneigh
  bytes += (double) nneigh_max * 3 * sizeof(double);              // rij"

  if (gradgradflag == 0)
    bytes += (double) nneigh_max * ndescriptors * 3 * sizeof(double);    // graddesc

  return bytes;
}