//
// CDDL HEADER START
//
// The contents of this file are subject to the terms of the Common Development
// and Distribution License Version 1.0 (the "License").
//
// You can obtain a copy of the license at
// http://www.opensource.org/licenses/CDDL-1.0.  See the License for the
// specific language governing permissions and limitations under the License.
//
// When distributing Covered Code, include this CDDL HEADER in each file and
// include the License file in a prominent location with the name LICENSE.CDDL.
// If applicable, add the following below this CDDL HEADER, with the fields
// enclosed by brackets "[]" replaced with your own identifying information:
//
// Portions Copyright (c) [yyyy] [name of copyright owner]. All rights reserved.
//
// CDDL HEADER END
//

//
// Copyright (c) 2018--2019, PANNAdevs group. 
// All rights reserved.
//
// Contributors:
//    Emine Kucukbenli
//    Franco Pellegrini
//    Ryan S. Elliott
//

#include "KIM_LogMacros.hpp"
#include "KIM_ModelDriverHeaders.hpp"
#include <algorithm>
#include <fstream>
#include <math.h>
#include <vector>

#define DIMENSION 3

#define SNUM(x)                                                \
  static_cast<std::ostringstream const &>(std::ostringstream() \
                                          << std::dec << x).str()


namespace
{
typedef double VectorOfSizeThree[3];

class PANNA
{
 public:
  //****************************************************************************
#undef KIM_LOGGER_OBJECT_NAME
#define KIM_LOGGER_OBJECT_NAME modelDriverCreate
  PANNA(KIM::ModelDriverCreate * const modelDriverCreate,
        KIM::LengthUnit const requestedLengthUnit,
        KIM::EnergyUnit const requestedEnergyUnit,
        KIM::ChargeUnit const requestedChargeUnit,
        KIM::TemperatureUnit const requestedTemperatureUnit,
        KIM::TimeUnit const requestedTimeUnit,
        int * const error) :
      modelWillNotRequestNeighborsOfNoncontributingParticles_(1)
  {
    int numberParameterFiles;
    modelDriverCreate->GetNumberOfParameterFiles(&numberParameterFiles);
    std::vector<std::string const *> paramFileNames(numberParameterFiles);
    for(int i=0;i<numberParameterFiles;i++){
      if(modelDriverCreate->GetParameterFileName(i,&(paramFileNames[i]))){
        LOG_ERROR("Could not get parameter file name.");
        *error = true;
        return;
      }
    }
    int gpout = get_parameters(modelDriverCreate, paramFileNames);
    if(gpout==0){
      LOG_DEBUG("Network loaded!");
    }
    else{
      LOG_ERROR("Error " + SNUM(gpout) + " while loading network!");
      *error = true;
      return;
    }

    *error = ConvertUnits(modelDriverCreate,
                          requestedLengthUnit,
                          requestedEnergyUnit,
                          requestedChargeUnit,
                          requestedTemperatureUnit,
                          requestedTimeUnit);
    if (*error) return;

    compute_parameters();

    modelDriverCreate->SetModelNumbering(KIM::NUMBERING::zeroBased);

    influenceDistance_ = par_.cutmax;
    cutoff_ = par_.cutmax;
    cutoffSq_ = cutoff_ * cutoff_;
    modelDriverCreate->SetInfluenceDistancePointer(&influenceDistance_);
    modelDriverCreate->SetNeighborListPointers(
        1, &cutoff_, &modelWillNotRequestNeighborsOfNoncontributingParticles_);

    for (int i=0; i<par_.Nspecies; ++i)
    {
      modelDriverCreate->SetSpeciesCode(KIM::SpeciesName(par_.species[i]), i+1);
    }

    // use function pointer declarations to verify prototypes
    KIM::ModelComputeArgumentsCreateFunction * CACreate
        = PANNA::ComputeArgumentsCreate;
    KIM::ModelComputeFunction * compute = PANNA::Compute;
    KIM::ModelComputeArgumentsDestroyFunction * CADestroy
        = PANNA::ComputeArgumentsDestroy;
    KIM::ModelDestroyFunction * destroy = PANNA::Destroy;

    *error = modelDriverCreate->SetRoutinePointer(
                 KIM::MODEL_ROUTINE_NAME::ComputeArgumentsCreate,
                 KIM::LANGUAGE_NAME::cpp,
                 true,
                 reinterpret_cast<KIM::Function *>(CACreate))
             || modelDriverCreate->SetRoutinePointer(
                 KIM::MODEL_ROUTINE_NAME::Compute,
                 KIM::LANGUAGE_NAME::cpp,
                 true,
                 reinterpret_cast<KIM::Function *>(compute))
             || modelDriverCreate->SetRoutinePointer(
                 KIM::MODEL_ROUTINE_NAME::ComputeArgumentsDestroy,
                 KIM::LANGUAGE_NAME::cpp,
                 true,
                 reinterpret_cast<KIM::Function *>(CADestroy))
             || modelDriverCreate->SetRoutinePointer(
                 KIM::MODEL_ROUTINE_NAME::Destroy,
                 KIM::LANGUAGE_NAME::cpp,
                 true,
                 reinterpret_cast<KIM::Function *>(destroy));
    if (*error) return;

    // everything is good
    *error = false;
    return;
  };

  //****************************************************************************
  ~PANNA() {};

  // Radial gvect contribution (and derivative part)
  double Gradial_d(double const rdiff, int const indr, double * const dtmp)
  {
    double const cent = rdiff - par_.Rsi_rad[indr];
    double const gauss = exp(-par_.eta_rad * cent * cent);
    double const fc = 0.5 * (1.0 + cos(rdiff * par_.iRc_rad));
    *dtmp = (par_.iRc_rad_half * sin(rdiff * par_.iRc_rad)
             + par_.twoeta_rad * fc * cent)
            * gauss / rdiff;
    return gauss * fc;
  }

  // Angular gvect contribution (and derivative part)
  double Gangular_d(double const rdiff1,
                    double const rdiff2,
                    double const cosijk,
                    int const Rsi,
                    int const Thi,
                    double * const dtmp)
  {
  	double costmp = cosijk;
    if (costmp > 0.999999999) costmp = 0.999999999;
    if (costmp < -0.999999999) costmp = -0.999999999;
    double const eps = 0.01;
    double sinijk = sqrt(1.0 - costmp * costmp + eps * pow(par_.Thi_sin[Thi], 2));
    double const iRij = 1.0 / rdiff1;
    double const iRik = 1.0 / rdiff2;
    double const Rcent = 0.5 * (rdiff1 + rdiff2) - par_.Rsi_ang[Rsi];
    double const fcrad
        = 0.5 * (1.0 + par_.Thi_cos[Thi] * cosijk + par_.Thi_sin[Thi] * sinijk);
    double const fcij = 0.5 * (1.0 + cos(rdiff1 * par_.iRc_ang));
    double const fcik = 0.5 * (1.0 + cos(rdiff2 * par_.iRc_ang));
    double const mod_norm = pow((1.0 + sqrt(1.0 + eps * pow(par_.Thi_sin[Thi], 2)))
    	* 0.5, par_.zeta);
    double const fact0
        = 2.0 * exp(-par_.eta_ang * Rcent * Rcent) * pow(fcrad, par_.zeta - 1) 
          / mod_norm;
    double const fact1 = fact0 * fcij * fcik;
    double const fact2
        = par_.zeta_half * fact1
          * (par_.Thi_cos[Thi] - par_.Thi_sin[Thi] * cosijk / sinijk);
    double const fact3 = par_.iRc_ang_half * fact0 * fcrad;
    double const G = fact1 * fcrad;
    dtmp[0] = -iRij
              * (par_.eta_ang * Rcent * G + fact2 * cosijk * iRij
                 + fact3 * fcik * sin(rdiff1 * par_.iRc_ang));
    dtmp[1] = fact2 * iRij * iRik;
    dtmp[2] = -iRik
              * (par_.eta_ang * Rcent * G + fact2 * cosijk * iRik
                 + fact3 * fcij * sin(rdiff2 * par_.iRc_ang));
    return G;
  }

  // Function computing gvect and its derivative
  void compute_gvect(int const ind1,
                     VectorOfSizeThree const * const x,
                     int const * const type,
                     int const * const neighs,
                     int const num_neigh,
                     std::vector<double> & G,
                     std::vector<double> & dGdx)
  {
    double const posx = x[ind1][0];
    double const posy = x[ind1][1];
    double const posz = x[ind1][2];
    // Elements to store neigh list for angular part
    int nan = 0;
    std::vector<int> ang_neigh(num_neigh);
    std::vector<int> ang_type(num_neigh);
    std::vector<double> dists(num_neigh);
    std::vector<double> diffx(num_neigh);
    std::vector<double> diffy(num_neigh);
    std::vector<double> diffz(num_neigh);
    //
    // Loop on neighbours, compute radial part, store quantities for angular
    for (int n = 0; n < num_neigh; n++)
    {
      int const nind = neighs[n];
      double const dx = x[nind][0] - posx;
      double const dy = x[nind][1] - posy;
      double const dz = x[nind][2] - posz;
      double const Rij = sqrt(dx * dx + dy * dy + dz * dz);
      if (Rij < par_.Rc_rad)
      {
        // Add all radial parts
        int const indsh = (type[nind] - 1) * par_.RsN_rad;
        for (int indr = 0; indr < par_.RsN_rad; indr++)
        {
          double dtmp;
          // Getting the simple G and derivative part
          G[indsh + indr] += Gradial_d(Rij, indr, &dtmp);
          // Filling all derivatives
          int const indsh2 = (indsh + indr) * (num_neigh + 1) * 3;
          double const derx = dtmp * dx;
          double const dery = dtmp * dy;
          double const derz = dtmp * dz;
          dGdx[indsh2 + num_neigh * 3] += derx;
          dGdx[indsh2 + num_neigh * 3 + 1] += dery;
          dGdx[indsh2 + num_neigh * 3 + 2] += derz;
          dGdx[indsh2 + n * 3] -= derx;
          dGdx[indsh2 + n * 3 + 1] -= dery;
          dGdx[indsh2 + n * 3 + 2] -= derz;
        }
      }
      // If within radial cutoff, store quantities
      if (Rij < par_.Rc_ang)
      {
        ang_neigh[nan] = n;
        ang_type[nan] = type[nind];
        dists[nan] = Rij;
        diffx[nan] = dx;
        diffy[nan] = dy;
        diffz[nan] = dz;
        nan++;
      }
    }

    // Loop on angular neighbours and fill angular part
    for (int n = 0; n < nan - 1; n++)
    {
      for (int m = n + 1; m < nan; m++)
      {
        // Compute cosine
        double const cos_ijk
            = (diffx[n] * diffx[m] + diffy[n] * diffy[m] + diffz[n] * diffz[m])
              / (dists[n] * dists[m]);
        // Gvect shift due to species
        int const indsh = par_.typsh[ang_type[n] - 1][ang_type[m] - 1];
        // Loop over all bins
        for (int Rsi = 0; Rsi < par_.RsN_ang; Rsi++)
        {
          for (int Thi = 0; Thi < par_.ThetasN; Thi++)
          {
            double dtmp[3];
            int const indsh2 = Rsi * par_.ThetasN + Thi;
            // Adding the G part and computing derivative
            G[indsh + indsh2]
                += Gangular_d(dists[n], dists[m], cos_ijk, Rsi, Thi, dtmp);
            // Computing the derivative contributions
            double const dgdxj = dtmp[0] * diffx[n] + dtmp[1] * diffx[m];
            double const dgdyj = dtmp[0] * diffy[n] + dtmp[1] * diffy[m];
            double const dgdzj = dtmp[0] * diffz[n] + dtmp[1] * diffz[m];
            double const dgdxk = dtmp[1] * diffx[n] + dtmp[2] * diffx[m];
            double const dgdyk = dtmp[1] * diffy[n] + dtmp[2] * diffy[m];
            double const dgdzk = dtmp[1] * diffz[n] + dtmp[2] * diffz[m];
            // Filling all the interested terms
            int const indsh3 = (indsh + indsh2) * (num_neigh + 1) * 3;
            dGdx[indsh3 + ang_neigh[n] * 3] += dgdxj;
            dGdx[indsh3 + ang_neigh[n] * 3 + 1] += dgdyj;
            dGdx[indsh3 + ang_neigh[n] * 3 + 2] += dgdzj;
            dGdx[indsh3 + ang_neigh[m] * 3] += dgdxk;
            dGdx[indsh3 + ang_neigh[m] * 3 + 1] += dgdyk;
            dGdx[indsh3 + ang_neigh[m] * 3 + 2] += dgdzk;
            dGdx[indsh3 + num_neigh * 3] -= dgdxj + dgdxk;
            dGdx[indsh3 + num_neigh * 3 + 1] -= dgdyj + dgdyk;
            dGdx[indsh3 + num_neigh * 3 + 2] -= dgdzj + dgdzk;
          }
        }
      }
    }
  }

  double compute_network(std::vector<double> & G,
                         std::vector<double> & dEdG,
                         int const type)
  {
    std::vector<double> lay1(G);
    std::vector<double> dlay1(par_.layers_size[type][0] * par_.gsize, 0.0);
    for (int i = 0; i < par_.gsize; i++) dlay1[i * par_.gsize + i] = 1.0;
    // Loop over layers
    for (int l = 0; l < par_.Nlayers[type]; l++)
    {
      int const size1 = par_.layers_size[type][l];
      int const size2 = par_.layers_size[type][l + 1];
      std::vector<double> lay2(size2);
      std::vector<double> dlay2(size2 * par_.gsize, 0.0);
      // Matrix vector multiplication done by hand for now
      // We compute W.x+b and W.(dx/dg)
      for (int i = 0; i < size2; i++)
      {
        // a_i = b_i
        lay2[i] = network_[type][2 * l + 1][i];
        for (int j = 0; j < size1; j++)
        {
          // a_i += w_ij * x_j
          lay2[i] += network_[type][2 * l][i * size1 + j] * lay1[j];
          for (int k = 0; k < par_.gsize; k++)
            // da_i/dg_k += w_ij * dx_j/dg_k
            dlay2[i * par_.gsize + k] += network_[type][2 * l][i * size1 + j]
                                         * dlay1[j * par_.gsize + k];
        }
      }
      // Apply appropriate activation
      // Gaussian
      if(par_.layers_activation[type][l]==1){
        for(int i=0; i<size2; i++){
          double const tmp = exp(-lay2[i]*lay2[i]);
          for(int k=0; k<par_.gsize; k++)
            dlay2[i*par_.gsize+k] *= -2.0*lay2[i]*tmp;
          lay2[i] = tmp;
        }
      }
      // ReLU
      else if(par_.layers_activation[type][l]==3){
        for(int i=0; i<size2; i++){
          if(lay2[i]<0.0){
            lay2[i] = 0.0;
            for(int k=0; k<par_.gsize; k++) dlay2[i*par_.gsize+k] = 0.0;
          }
        }
      }
      // Otherwise it's linear and nothing needs to be done

      lay1 = lay2;
      dlay1 = dlay2;
    }
    for(int k=0; k<par_.gsize; k++) dlay1[k] *= par_.energy_conversion;
    dEdG = dlay1;
    double const E = lay1[0]*par_.energy_conversion;
    return E;
  }

  // ########################################################
  // ########################################################

  // Get a new line skipping comments or empty lines
  // Set value=... if [...], return 1
  // Fill key,value if 'key=value', return 2
  // Set value=... if ..., return 3
  // Return 0 if eof, <0 if error, >0 if okay
  int get_input_line(KIM::ModelDriverCreate * const modelDriverCreate,
                     std::ifstream* file, std::string* key, std::string* value){
    std::string line;
    int parsed = 0;
    while(!parsed){
      std::getline(*file,line);
      // Exit on EOF
      if(file->eof()) return 0;
      // Exit on bad read
      if(file->bad()) return -1;
      // Remove spaces
      line.erase (std::remove(line.begin(), line.end(), ' '), line.end());
      // Skip empty line
      if(line.length()==0) continue;
      // Skip comments
      if(line.at(0)=='#') continue;
      // Parse headers
      if(line.at(0)=='['){
        *value = line.substr(1,line.length()-2);
        return 1;
      }
      // Look for equal sign
      std::string eq = "=";
      size_t eqpos = line.find(eq);
      // Parse key-value pair
      if(eqpos != std::string::npos){
        *key = line.substr(0,eqpos);
        *value = line.substr(eqpos+1,line.length()-1);
        return 2;
      }
      // Parse full line
      else{
        *value = line;
        return 3;
      }

      LOG_DEBUG(line);
      parsed = 1;
    }
    return -1;
  }

  // int get_parameters(char* directory, char* filename)
  int get_parameters(KIM::ModelDriverCreate * const modelDriverCreate,
                     std::vector<std::string const *> paramFileNames)
  {
    // Parsing the potential parameters
    std::ifstream params_file;
    std::ifstream weights_file;
    std::string key, value;

    // Initializing some parameters before reading:
    par_.Nspecies = -1;
    // Flags to keep track of set parameters
    int Npars = 14;
    std::vector<int> parset(Npars,0);
    std::vector<int> spset;

    params_file.open(paramFileNames[0]->c_str());
    // section keeps track of input file sections
    // -1 in the beginning
    // 0 for gvect params
    // i for species i (1 based)
    int section = -1;
    // parseint checks the status of input parsing
    int parseint = get_input_line(modelDriverCreate,&params_file,&key,&value);
    while(parseint>0){
      // Parse line
      if(parseint==1){
        // Gvect param section
        if(value=="GVECT_PARAMETERS"){
          section = 0;
        }
        // For now other sections are just species networks
        else {
          // First time after params: do checks
          if(section==0){
            // Set steps if they were omitted
            if(parset[5]==0){
              par_.Rsst_rad = (par_.Rc_rad - par_.Rs0_rad) / par_.RsN_rad;
              parset[5]=1;
            }
            if(parset[10]==0){
              par_.Rsst_ang = (par_.Rc_ang - par_.Rs0_ang) / par_.RsN_ang;
              parset[10]=1;
            }
            // Check that all parameters have been set
            for(int p=0;p<Npars;p++){
              if(parset[p]==0){
                LOG_DEBUG("Parameter " + SNUM(p) + " not set!");
                return -1;
              }
            }
            // Calculate Gsize
            par_.gsize = par_.Nspecies * par_.RsN_rad + (par_.Nspecies*(par_.Nspecies+1))/2 * par_.RsN_ang * par_.ThetasN;
          }
          int match = 0;
          for(int s=0;s<par_.Nspecies;s++){
            // If species matches the list, change section
            if(value==par_.species[s]){
              section = s+1;
              match = 1;
            }
          }
          if(match==0){
            LOG_DEBUG("Species " + value + " not found in species list.");
            return -2;
          }
        }
      }
      else if(parseint==2){
        // Parse param section
        if(section==0){
          if(key=="Nspecies"){
            par_.Nspecies = std::atoi(value.c_str());
            // Small check
            if(par_.Nspecies<1){
              LOG_DEBUG("Nspecies needs to be >0.");
              return -2;
            }
            if(par_.Nspecies!=((int)(paramFileNames.size())-1)){
              LOG_DEBUG("Nspecies does not match the number of files.");
              return -2;
            }
            parset[0] = 1;
            par_.species = std::vector<std::string>(par_.Nspecies);
            par_.Nlayers = std::vector<int>(par_.Nspecies,-1);
            par_.layers_size = std::vector<std::vector<int> >(par_.Nspecies);
            par_.layers_activation = std::vector<std::vector<int> >(par_.Nspecies);
            network_ = std::vector<std::vector<std::vector<double> > >(par_.Nspecies);
            // Keep track of set species
            spset = std::vector<int>(par_.Nspecies,0);
          }
          else if(key=="species"){
            std::string comma = ",";
            size_t pos = 0;
            int s = 0;
            // Parse species list
            while ((pos = value.find(comma)) != std::string::npos) {
              if(s>par_.Nspecies-2){
                LOG_DEBUG("Species list longer than Nspecies.");
                return -2;
              }
              par_.species[s] = value.substr(0, pos);
              value.erase(0, pos+1);
              s++;
            }
            if(value.length()>0){
              par_.species[s] = value;
              s++;
            };
            if(s<par_.Nspecies){
              LOG_DEBUG("Species list shorter than Nspecies.");
              return -2;
            }
            parset[1] = 1;
          }
          else if(key=="eta_rad"){
            par_.eta_rad = std::atof(value.c_str());
            parset[2] = 1;
          }
          else if(key=="Rc_rad"){
            par_.Rc_rad = std::atof(value.c_str());
            parset[3] = 1;
          }
          else if(key=="Rs0_rad"){
            par_.Rs0_rad = std::atof(value.c_str());
            parset[4] = 1;
          }
          else if(key=="Rsst_rad"){
            par_.Rsst_rad = std::atof(value.c_str());
            parset[5] = 1;
          }
          else if(key=="RsN_rad"){
            par_.RsN_rad = std::atoi(value.c_str());
            parset[6] = 1;
          }
          else if(key=="eta_ang"){
            par_.eta_ang = std::atof(value.c_str());
            parset[7] = 1;
          }
          else if(key=="Rc_ang"){
            par_.Rc_ang = std::atof(value.c_str());
            parset[8] = 1;
          }
          else if(key=="Rs0_ang"){
            par_.Rs0_ang = std::atof(value.c_str());
            parset[9] = 1;
          }
          else if(key=="Rsst_ang"){
            par_.Rsst_ang = std::atof(value.c_str());
            parset[10] = 1;
          }
          else if(key=="RsN_ang"){
            par_.RsN_ang = std::atoi(value.c_str());
            parset[11] = 1;
          }
          else if(key=="zeta"){
            par_.zeta = std::atof(value.c_str());
            parset[12] = 1;
          }
          else if(key=="ThetasN"){
            par_.ThetasN = std::atoi(value.c_str());
            parset[13] = 1;
          }
        }
        // Parse species network
        else if(section<par_.Nspecies+1){
          int s=section-1;
          // Read species network
          if(key=="Nlayers"){
            par_.Nlayers[s] = std::atoi(value.c_str());
            // This has the extra gvect size
            par_.layers_size[s] = std::vector<int>(par_.Nlayers[s]+1);
            par_.layers_size[s][0] = par_.gsize;
            par_.layers_size[s][1] = 0;
            par_.layers_activation[s] = std::vector<int>(par_.Nlayers[s],1);
            par_.layers_activation[s][par_.Nlayers[s]-1]=0;
            network_[s] = std::vector<std::vector<double> >(2*par_.Nlayers[s]);
          }
          else if(key=="sizes"){
            if(par_.Nlayers[s]==-1){
              LOG_DEBUG("Sizes cannot be set before Nlayers.");
              return -3;
            }
            std::string comma = ",";
            size_t pos = 0;
            int l = 0;
            // Parse layers list
            while ((pos = value.find(comma)) != std::string::npos) {
              if(l>par_.Nlayers[s]-2){
                LOG_DEBUG("Layers list longer than Nlayers.");
                return -3;
              }
              std::string lsize = value.substr(0, pos);
              par_.layers_size[s][l+1] = std::atoi(lsize.c_str());
              value.erase(0, pos+1);
              l++;
            }
            if(value.length()>0){
              par_.layers_size[s][l+1] = std::atoi(value.c_str());
              l++;
            };
            if(l<par_.Nlayers[s]){
              LOG_DEBUG("Layers list shorter than Nlayers.");
              return -3;
            }
          }
          else if(key=="activations"){
            if(par_.Nlayers[s]==-1){
              LOG_DEBUG("Activations cannot be set before Nlayers.");
              return -3;
            }
            std::string comma = ",";
            size_t pos = 0;
            int l = 0;
            // Parse layers list
            while ((pos = value.find(comma)) != std::string::npos) {
              if(l>par_.Nlayers[s]-2){
                LOG_DEBUG("Activations list longer than Nlayers.");
                return -3;
              }
              std::string lact = value.substr(0, pos);
              int actnum = std::atoi(lact.c_str());
              if (actnum!=0 && actnum!=1 && actnum!=3){
                LOG_DEBUG("Activations unsupported: " + SNUM(actnum));
                return -3;
              }
              par_.layers_activation[s][l] = actnum;
              value.erase(0, pos+1);
              l++;
            }
            if(value.length()>0){
              int actnum = std::atoi(value.c_str());
              if (actnum!=0 && actnum!=1 && actnum!=3){
                LOG_DEBUG("Activations unsupported: " + SNUM(actnum));
                return -3;
              }
              par_.layers_activation[s][l] = actnum;
              l++;
            };
            if(l<par_.Nlayers[s]){
              LOG_DEBUG("Activations list shorter than Nlayers.");
              return -3;
            }
          }
          else if(key=="file"){
            if(par_.layers_size[s][1]==0){
              LOG_DEBUG("Layers sizes unset before filename for species "
                        + par_.species[s]);
              return -3;
            }
            // Read file and load weights
            weights_file.open(paramFileNames[section]->c_str(), std::ios::binary);
            if(!weights_file.is_open()){
              LOG_DEBUG("Error reading weights file for " + par_.species[s]);
              return -3;
            }
            for(int l=0; l<par_.Nlayers[s]; l++){
              // Allocate and read the right amount of data
              // Weights
              network_[s][2*l] = std::vector<double>(par_.layers_size[s][l]*par_.layers_size[s][l+1]);
              for(int i=0; i<par_.layers_size[s][l]; i++) {
                for(int j=0; j<par_.layers_size[s][l+1]; j++) {
                  float num;
                  weights_file.read(reinterpret_cast<char*>(&num), sizeof(float));
                  if(weights_file.eof()){
                    LOG_DEBUG("Weights file for " + par_.species[s]
                              + " is too small.");
                    return -3;
                  }
                  network_[s][2*l][j*par_.layers_size[s][l]+i] = (double)num;
                }
              }
              // Biases
              network_[s][2*l+1] = std::vector<double>(par_.layers_size[s][l+1]);
              for(int d=0; d<par_.layers_size[s][l+1]; d++) {
                float num;
                weights_file.read(reinterpret_cast<char*>(&num), sizeof(float));
                if(weights_file.eof()){
                  LOG_DEBUG("Weights file for " + par_.species[s]
                            + " is too small.");
                  return -3;
                }
                network_[s][2*l+1][d] = (double)num;
              }
            }
            // Check if we're not at the end
            std::ifstream::pos_type fpos = weights_file.tellg();
            weights_file.seekg(0, std::ios::end);
            std::ifstream::pos_type epos = weights_file.tellg();
            if(fpos!=epos){
              LOG_DEBUG("Weights file for " + par_.species[s] + " is too big.");
              return -3;
            }
            weights_file.close();
            spset[section-1] = 1;
          }
        }
        else{
          return -3;
        }
      }
      else if(parseint==3){
        // No full line should be in the input
        LOG_DEBUG("Unexpected line " + value);
        return -4;
      }

      // Get new line
      parseint = get_input_line(modelDriverCreate,&params_file,&key,&value);
    }
    params_file.close();
    return(0);
}

void compute_parameters(){
    // Compute derived params
    par_.cutmax = par_.Rc_rad>par_.Rc_ang ? par_.Rc_rad : par_.Rc_ang;
    par_.seta_rad = sqrt(par_.eta_rad);
    par_.twoeta_rad = 2.0*par_.eta_rad;
    par_.seta_ang = sqrt(par_.eta_ang);
    par_.zint = (int) par_.zeta;
    par_.zeta_half = 0.5*par_.zeta;
    par_.iRc_rad = M_PI/par_.Rc_rad;
    par_.iRc_rad_half = 0.5*par_.iRc_rad;
    par_.iRc_ang = M_PI/par_.Rc_ang;
    par_.iRc_ang_half = 0.5*par_.iRc_ang;
    for (int indr = 0; indr < par_.RsN_rad; indr++)
      par_.Rsi_rad.push_back(par_.Rs0_rad + indr * par_.Rsst_rad);
    for (int indr = 0; indr < par_.RsN_ang; indr++)
      par_.Rsi_ang.push_back(par_.Rs0_ang + indr * par_.Rsst_ang);
    for (int indr = 0; indr < par_.ThetasN; indr++)
    {
      double ti = (indr + 0.5f) * M_PI / par_.ThetasN;
      par_.Thi_cos.push_back(cos(ti));
      par_.Thi_sin.push_back(sin(ti));
    }

    // Precalculate gvect shifts for any species pair
    for (int s = 0; s < par_.Nspecies; s++)
    {
      par_.typsh.push_back(std::vector<int>(par_.Nspecies));
      for (int ss = 0; ss < par_.Nspecies; ss++)
      {
        if (s < ss)
          par_.typsh[s][ss] = par_.Nspecies * par_.RsN_rad
                              + (s * par_.Nspecies - (s * (s + 1)) / 2 + ss)
                                    * par_.RsN_ang * par_.ThetasN;
        else
          par_.typsh[s][ss] = par_.Nspecies * par_.RsN_rad
                              + (ss * par_.Nspecies - (ss * (ss + 1)) / 2 + s)
                                    * par_.RsN_ang * par_.ThetasN;
      }
    }
  }

  //****************************************************************************
  // no need to make these "extern" since KIM will only access them
  // via function pointers.  "static" is required so that there is not
  // an implicit this pointer added to the prototype by the C++ compiler
  static int Destroy(KIM::ModelDestroy * const modelDestroy)
  {
    PANNA * model;
    modelDestroy->GetModelBufferPointer(reinterpret_cast<void **>(&model));

    if (model != NULL)
    {
      // delete object itself
      delete model;
    }

    // everything is good
    return false;
  }

  //****************************************************************************
#undef KIM_LOGGER_OBJECT_NAME
#define KIM_LOGGER_OBJECT_NAME modelCompute
  //
  static int
  Compute(KIM::ModelCompute const * const modelCompute,
          KIM::ModelComputeArguments const * const modelComputeArguments)
  {
    int const * numberOfParticlesPointer;
    int const * particleSpeciesCodes;
    int const * particleContributing;
    double const * coordinates;
    double * partialEnergy;
    double * partialParticleEnergy;
    double * partialForces;

    PANNA * panna;
    modelCompute->GetModelBufferPointer(reinterpret_cast<void **>(&panna));

    int error = modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::numberOfParticles,
                    &numberOfParticlesPointer)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::particleSpeciesCodes,
                    &particleSpeciesCodes)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::particleContributing,
                    &particleContributing)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::coordinates,
                    (double const **) &coordinates)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::partialEnergy, &partialEnergy)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::partialParticleEnergy,
                    &partialParticleEnergy)
                || modelComputeArguments->GetArgumentPointer(
                    KIM::COMPUTE_ARGUMENT_NAME::partialForces,
                    (double const **) &partialForces);
    if (error)
    {
      LOG_ERROR("Unable to get argument pointers");
      return error;
    }

    int const inum = *numberOfParticlesPointer;
    VectorOfSizeThree const * const x = (VectorOfSizeThree const *) coordinates;
    VectorOfSizeThree * const f = (VectorOfSizeThree *) partialForces;
    int const * type = particleSpeciesCodes;

    // initialize energy and forces
    if (partialEnergy) *partialEnergy = 0.0;
    if (partialParticleEnergy)
      for (int i = 0; i < inum; ++i) { partialParticleEnergy[i] = 0.0; }
    int const extent = inum * DIMENSION;
    if (partialForces)
      for (int i = 0; i < extent; ++i) { partialForces[i] = 0.0; }

    int numberOfNeighbors;
    int const * neighbors;

    // Allocate this gvect and dG/dx
    std::vector<double> G(panna->par_.gsize);
    std::vector<double> dEdG(panna->par_.gsize);

    // Looping on contributing atoms
    for (int myind = 0; myind < inum; myind++)
    {
      if (particleContributing[myind])
      {
        modelComputeArguments->GetNeighborList(
            0, myind, &numberOfNeighbors, &neighbors);

        // dGdx has (numn+1)*3 derivs per elem: neigh first, then the atom
        // itself
        std::vector<double> dGdx(panna->par_.gsize * (numberOfNeighbors + 1)
                                 * 3);
        for (int i = 0; i < panna->par_.gsize; i++)
        {
          G[i] = 0.0;
          for (int j = 0; j < (numberOfNeighbors + 1) * 3; j++)
            dGdx[i * (numberOfNeighbors + 1) * 3 + j] = 0.0;
        }
        // Calculate Gvect and derivatives
        panna->compute_gvect(
            myind, x, type, neighbors, numberOfNeighbors, G, dGdx);

        // Apply network
        double E = panna->compute_network(G, dEdG, type[myind] - 1);
        // Calculate forces
        int shift = (numberOfNeighbors + 1) * 3;
        if (f)
        {
          for (int n = 0; n < numberOfNeighbors; n++)
          {
            int nind = neighbors[n];
            for (int j = 0; j < panna->par_.gsize; j++)
            {
              f[nind][0] -= dEdG[j] * dGdx[j * shift + 3 * n];
              f[nind][1] -= dEdG[j] * dGdx[j * shift + 3 * n + 1];
              f[nind][2] -= dEdG[j] * dGdx[j * shift + 3 * n + 2];
            }
          }
          for (int j = 0; j < panna->par_.gsize; j++)
          {
            f[myind][0] -= dEdG[j] * dGdx[j * shift + 3 * numberOfNeighbors];
            f[myind][1] -= dEdG[j] * dGdx[j * shift + 3 * numberOfNeighbors + 1];
            f[myind][2] -= dEdG[j] * dGdx[j * shift + 3 * numberOfNeighbors + 2];
          }
        }

        if (partialEnergy) *partialEnergy += E;
        if (partialParticleEnergy) partialParticleEnergy[myind] = E;
      }  // loop on myind
    }

    // everything is good
    return false;
  };

  //****************************************************************************
  static int ComputeArgumentsCreate(
      KIM::ModelCompute const * const /* modelCompute */,
      KIM::ModelComputeArgumentsCreate * const modelComputeArgumentsCreate)
  {
    // register arguments
    int error = modelComputeArgumentsCreate->SetArgumentSupportStatus(
                    KIM::COMPUTE_ARGUMENT_NAME::partialEnergy,
                    KIM::SUPPORT_STATUS::optional)
                || modelComputeArgumentsCreate->SetArgumentSupportStatus(
                    KIM::COMPUTE_ARGUMENT_NAME::partialParticleEnergy,
                    KIM::SUPPORT_STATUS::optional)
                || modelComputeArgumentsCreate->SetArgumentSupportStatus(
                    KIM::COMPUTE_ARGUMENT_NAME::partialForces,
                    KIM::SUPPORT_STATUS::optional);


    // register callbacks
    //
    // none

    return error;
  }

  //****************************************************************************
  static int
  ComputeArgumentsDestroy(KIM::ModelCompute const * const /* modelCompute */,
                          KIM::ModelComputeArgumentsDestroy * const
                          /* modelComputeArgumentsDestroy */)
  {
    // nothing further to do

    return false;
  }

 private:
  //****************************************************************************
  // Member variables
  double influenceDistance_;
  double cutoff_;
  double cutoffSq_;
  int const modelWillNotRequestNeighborsOfNoncontributingParticles_;
  struct parameters  // Gvect and NN parameters
  {
    int Nspecies;
    // Gvector parameters
    double eta_rad;
    double Rc_rad;
    double Rs0_rad;
    double Rsst_rad;
    int RsN_rad;
    double eta_ang;
    double Rc_ang;
    double Rs0_ang;
    double Rsst_ang;
    int RsN_ang;
    double zeta;
    int ThetasN;
    std::vector<std::string> species;
    int gsize;
    double energy_conversion;

    // Network parameters
    std::vector<int> Nlayers;
    std::vector<std::vector<int> > layers_size;
    std::vector<std::vector<int> > layers_activation;

    // Useful precalculated quantities
    double cutmax;
    double seta_rad;
    double twoeta_rad;
    double seta_ang;
    int zint;
    double zeta_half;
    double iRc_rad;
    double iRc_rad_half;
    double iRc_ang;
    double iRc_ang_half;
    std::vector<double> Rsi_rad;
    std::vector<double> Rsi_ang;
    std::vector<double> Thi_cos;
    std::vector<double> Thi_sin;
    std::vector<std::vector<int> > typsh;
  } par_;
  // The network [species, layers, array]
  std::vector<std::vector<std::vector<double> > > network_;

  //****************************************************************************
#undef KIM_LOGGER_OBJECT_NAME
#define KIM_LOGGER_OBJECT_NAME modelDriverCreate
  //
  int ConvertUnits(KIM::ModelDriverCreate * const modelDriverCreate,
                   KIM::LengthUnit const requestedLengthUnit,
                   KIM::EnergyUnit const requestedEnergyUnit,
                   KIM::ChargeUnit const requestedChargeUnit,
                   KIM::TemperatureUnit const requestedTemperatureUnit,
                   KIM::TimeUnit const requestedTimeUnit)
  {
    int ier;

    // define default base units
    KIM::LengthUnit fromLength = KIM::LENGTH_UNIT::A;
    KIM::EnergyUnit fromEnergy = KIM::ENERGY_UNIT::eV;
    KIM::ChargeUnit fromCharge = KIM::CHARGE_UNIT::unused;
    KIM::TemperatureUnit fromTemperature = KIM::TEMPERATURE_UNIT::unused;
    KIM::TimeUnit fromTime = KIM::TIME_UNIT::unused;

    // changing units
    double convertLength = 1.0;
    ier = KIM::ModelDriverCreate::ConvertUnit(fromLength,
                                              fromEnergy,
                                              fromCharge,
                                              fromTemperature,
                                              fromTime,
                                              requestedLengthUnit,
                                              requestedEnergyUnit,
                                              requestedChargeUnit,
                                              requestedTemperatureUnit,
                                              requestedTimeUnit,
                                              1.0,
                                              0.0,
                                              0.0,
                                              0.0,
                                              0.0,
                                              &convertLength);
    if (ier)
    {
      LOG_ERROR("Unable to convert length unit");
      return ier;
    }

    par_.eta_rad /= convertLength*convertLength;
    par_.Rc_rad *= convertLength;
    par_.Rs0_rad *= convertLength;
    par_.Rsst_rad *= convertLength;
    par_.eta_ang /= convertLength*convertLength;
    par_.Rc_ang *= convertLength;
    par_.Rs0_ang *= convertLength;
    par_.Rsst_ang *= convertLength;

    par_.energy_conversion = 1.0;
    ier = KIM::ModelDriverCreate::ConvertUnit(fromLength,
                                              fromEnergy,
                                              fromCharge,
                                              fromTemperature,
                                              fromTime,
                                              requestedLengthUnit,
                                              requestedEnergyUnit,
                                              requestedChargeUnit,
                                              requestedTemperatureUnit,
                                              requestedTimeUnit,
                                              0.0,
                                              1.0,
                                              0.0,
                                              0.0,
                                              0.0,
                                              &par_.energy_conversion);
    if (ier)
    {
      LOG_ERROR("Unable to convert energy unit");
      return ier;
    }

    // register units
    ier = modelDriverCreate->SetUnits(requestedLengthUnit,
                                      requestedEnergyUnit,
                                      KIM::CHARGE_UNIT::unused,
                                      KIM::TEMPERATURE_UNIT::unused,
                                      KIM::TIME_UNIT::unused);
    if (ier)
    {
      LOG_ERROR("Unable to set units to requested values");
      return ier;
    }

    // everything is good
    ier = false;
    return ier;
  }
};

}  // namespace

//******************************************************************************
extern "C" {
int model_driver_create(KIM::ModelDriverCreate * const modelDriverCreate,
                        KIM::LengthUnit const requestedLengthUnit,
                        KIM::EnergyUnit const requestedEnergyUnit,
                        KIM::ChargeUnit const requestedChargeUnit,
                        KIM::TemperatureUnit const requestedTemperatureUnit,
                        KIM::TimeUnit const requestedTimeUnit)
{
  int ier;

  // read input files, convert units if needed, compute
  // interpolation coefficients, set cutoff, and publish parameters
  PANNA * const modelObject = new PANNA(modelDriverCreate,
                                        requestedLengthUnit,
                                        requestedEnergyUnit,
                                        requestedChargeUnit,
                                        requestedTemperatureUnit,
                                        requestedTimeUnit,
                                        &ier);
  if (ier)
  {
    // constructor already reported the error
    delete modelObject;
    return ier;
  }

  // register pointer to PANNA object in KIM object
  modelDriverCreate->SetModelBufferPointer(static_cast<void *>(modelObject));

  // everything is good
  ier = false;
  return ier;
}
}  // extern "C"