/* ----------------------------------------------------------------------   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   http://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 authors: German Samolyuk (ORNL) and
                         Mario Pinto (Computational Research Lab, Pune, India)
------------------------------------------------------------------------- */

#include <math.h>
#include <string.h>
#include "compute_gkma.h"
#include "atom.h"
#include "update.h"
#include "modify.h"
#include "force.h"
#include "group.h"
#include "error.h"
#include "memory.h" //GKMA
#include <iostream> //GKMA
#include <fstream> //GKMA
#include <string> //GKMA

using namespace LAMMPS_NS;

#define INVOKED_PERATOM 8

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

ComputeGKMA::ComputeGKMA(LAMMPS *lmp, int narg, char **arg) :
  Compute(lmp, narg, arg),
  id_ke(NULL), id_pe(NULL), id_stress(NULL)
{
  if (narg != 9) error->all(FLERR,"Illegal compute GKMA command");

  //GKMA-beg-1
  bigint natoms = atom->natoms;
  array_flag = 1;
  firstmode = force->inumeric(FLERR,arg[6]);
  lastmode = force->inumeric(FLERR,arg[7]);
  binsize = force->inumeric(FLERR,arg[8]);
  nummodes = lastmode-firstmode+1;
  numbins = floor(nummodes/binsize);
  size_array_rows = numbins;
  size_array_cols = 6;
  extarray = 1;
  //GKMA-end-1

  // store ke/atom, pe/atom, stress/atom IDs used by heat flux computation
  // insure they are valid for these computations

  int n = strlen(arg[3]) + 1;
  id_ke = new char[n];
  strcpy(id_ke,arg[3]);

  n = strlen(arg[4]) + 1;
  id_pe = new char[n];
  strcpy(id_pe,arg[4]);

  n = strlen(arg[5]) + 1;
  id_stress = new char[n];
  strcpy(id_stress,arg[5]);

  int ike = modify->find_compute(id_ke);
  int ipe = modify->find_compute(id_pe);
  int istress = modify->find_compute(id_stress);
  if (ike < 0 || ipe < 0 || istress < 0)
    error->all(FLERR,"Could not find compute gkma compute ID");
  if (strcmp(modify->compute[ike]->style,"ke/atom") != 0)
    error->all(FLERR,"Compute gkma compute ID does not compute ke/atom");
  if (modify->compute[ipe]->peatomflag == 0)
    error->all(FLERR,"Compute gkmka compute ID does not compute pe/atom");
  if (modify->compute[istress]->pressatomflag == 0)
    error->all(FLERR,
               "Compute gkma compute ID does not compute stress/atom");

  memory->create(array,3*natoms,6,"gkma:array");

  //GKMA-beg-2
  std::ifstream eigfile;
  eigfile.open("eignvector.eig");

  std::string val;
  double doubleval;
  memory->create(eigx,nummodes*natoms,"gkma:eigx");
  memory->create(eigy,nummodes*natoms,"gkma:eigy");
  memory->create(eigz,nummodes*natoms,"gkma:eigz");
  
  for (int i=0; i<=natoms+3 ; i++){
    getline(eigfile,val);
  }
  for (int i=0; i<firstmode-1; i++){
    for (int j=0; j<natoms+2; j++) getline(eigfile,val);
  }
  for (int i=0; i<nummodes; i++){
    getline(eigfile,val);
    getline(eigfile,val);
    for (int j=0; j<natoms; j++){
      eigfile >> doubleval;
      eigx[i*natoms+j]=doubleval; //values are grouped by eigenvector, not by atom
      eigfile >> doubleval;
      eigy[i*natoms+j]=doubleval;
      eigfile >> doubleval;
      eigz[i*natoms+j]=doubleval;
    }
    getline(eigfile,val);
  }
  eigfile.close();
  //GKMA-end-2
}

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

ComputeGKMA::~ComputeGKMA()
{
  delete [] id_ke;
  delete [] id_pe;
  delete [] id_stress;
  delete [] array;  //GKMA
}

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

void ComputeGKMA::init()
{
  // error checks

  int ike = modify->find_compute(id_ke);
  int ipe = modify->find_compute(id_pe);
  int istress = modify->find_compute(id_stress);
  if (ike < 0 || ipe < 0 || istress < 0)
    error->all(FLERR,"Could not find compute gkma compute ID");

  c_ke = modify->compute[ike];
  c_pe = modify->compute[ipe];
  c_stress = modify->compute[istress];
}

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

void ComputeGKMA::compute_array()
{
  invoked_array = update->ntimestep;
  
  // invoke 3 computes if they haven't been already

  if (!(c_ke->invoked_flag & INVOKED_PERATOM)) {
    c_ke->compute_peratom();
    c_ke->invoked_flag |= INVOKED_PERATOM;
  }
  if (!(c_pe->invoked_flag & INVOKED_PERATOM)) {
    c_pe->compute_peratom();
    c_pe->invoked_flag |= INVOKED_PERATOM;
  }
  if (!(c_stress->invoked_flag & INVOKED_PERATOM)) {
    c_stress->compute_peratom();
    c_stress->invoked_flag |= INVOKED_PERATOM;
  }

  // heat flux vector = jc[3] + jv[3]
  // jc[3] = convective portion of heat flux = sum_i (ke_i + pe_i) v_i[3]
  // jv[3] = virial portion of heat flux = sum_i (stress_tensor_i . v_i[3])
  // normalization by volume is not included

  double *ke = c_ke->vector_atom;
  double *pe = c_pe->vector_atom;
  double **stress = c_stress->array_atom;

  double **v = atom->v;
  int *mask = atom->mask;
  int nlocal = atom->nlocal;
  //GKMA-beg-3
  bigint natoms = atom->natoms;
  double **f = atom->f;
  int *type = atom->type;
  double *mass = atom->mass;
  tagint *tag = atom->tag;

  double jcx[nummodes];
  double jcy[nummodes];
  double jcz[nummodes];
  double jvx[nummodes];
  double jvy[nummodes];
  double jvz[nummodes];
  for (int i=0; i < nummodes; i++) {
    jcx[i]=jcy[i]=jcz[i]=jvx[i]=jvy[i]=jvz[i]=0;
  }
  double eng;
  double xdotx[nummodes];
  double xdoty[nummodes];
  double xdotz[nummodes];
  double vx;
  double vy;
  double vz;

  double vsum[nlocal], actv[nlocal];

  for (int i = 0; i < nummodes; i++) {
    xdotx[i]=xdoty[i]=xdotz[i]=0;
    for (int j = 0; j < nlocal; j++) {
      if (mask[j] & groupbit) {
        xdotx[i] += sqrt(mass[type[j]])*eigx[i*natoms+tag[j]-1]*v[j][0];
        xdoty[i] += sqrt(mass[type[j]])*eigy[i*natoms+tag[j]-1]*v[j][1];
        xdotz[i] += sqrt(mass[type[j]])*eigz[i*natoms+tag[j]-1]*v[j][2];
      }
    }
  }
  for (int i = 0; i < nlocal; i++) {
    if (mask[i] & groupbit) {
      eng = pe[i] + ke[i];
      vsum[i] = actv[i] = 0;
      for (int j = 0; j < nummodes; j++) {
        vx=1/sqrt(mass[type[i]])*eigx[j*natoms+tag[i]-1]*xdotx[j];
        vy=1/sqrt(mass[type[i]])*eigy[j*natoms+tag[i]-1]*xdoty[j];
        vz=1/sqrt(mass[type[i]])*eigz[j*natoms+tag[i]-1]*xdotz[j];
        jcx[j] += eng*vx;
        jcy[j] += eng*vy;
        jcz[j] += eng*vz;
        jvx[j] -= stress[i][0]*vx + stress[i][3]*vy +
          stress[i][4]*vz;
        jvy[j] -= stress[i][3]*vx + stress[i][1]*vy +
          stress[i][5]*vz;
        jvz[j] -= stress[i][4]*vx + stress[i][5]*vy +
          stress[i][2]*vz;
        vsum[i] += vx;
        actv[i] = v[i][0];
      }
    }
  }

  // convert jv from stress*volume to energy units via nktv2p factor

  double data[numbins][6];
  for (int i = 0; i < numbins; i++) {
    data[i][0]=data[i][1]=data[i][2]=data[i][3]=data[i][4]=data[i][5]=0;
  }
  double nktv2p = force->nktv2p;
  for (int i = 0; i < numbins; i++) {
    for (int j = 0; j < binsize; j++) {
      data[i][0]+=jcx[j+i*binsize]+jvx[j+i*binsize]/nktv2p;
      data[i][1]+=jcy[j+i*binsize]+jvy[j+i*binsize]/nktv2p;
      data[i][2]+=jcz[j+i*binsize]+jvz[j+i*binsize]/nktv2p;
      data[i][3]+=jcx[j+i*binsize];
      data[i][4]+=jcy[j+i*binsize];
      data[i][5]+=jcz[j+i*binsize];
    }
  }
  
  // sum across all procs
  // 1st 3 terms are total heat flux
  // 2nd 3 terms are just conductive portion
  MPI_Allreduce(data[0],array[0],6*numbins,MPI_DOUBLE,MPI_SUM,world);
  //GKMA-end-3
}