/* ----------------------------------------------------------------------
   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 author: Oscar Samuel Cajahuaringa Macollunco sammu09@...4527...l.com
------------------------------------------------------------------------- */

#include "math.h"
#include "string.h"
#include "stdlib.h"
#include "stdio.h"
#include "fix_tasm.h"
#include "atom.h"
#include "update.h"
#include "group.h"
#include "modify.h"
#include "neighbor.h"
#include "force.h"
#include "pair.h"
#include "comm.h"
#include "compute.h"
#include "fix.h"
#include "domain.h"
#include "memory.h"
#include "timer.h"
#include "error.h"
#include "output.h"

using namespace LAMMPS_NS;
using namespace FixConst;

enum{X,Y,Z};

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

FixTASM::FixTASM(LAMMPS *lmp, int narg, char **arg) :
  Fix(lmp, narg, arg)
{  
  if (narg < 7 || narg > 11) 
    error->all(FLERR,"Illegal fix tasm command");

  nevery = force->inumeric(FLERR,arg[3]);

  if (nevery < 0) error->all(FLERR,"Illegal fix tasm command");

  scalar_flag = 0;
  vector_flag = 1;
  size_vector = 5;
  box_change_size = 1;
  box_change_shape = 0;
  extscalar = 1;
  extvector = 0;
  allremap = 1;
 
  temp_tasm = force->numeric(FLERR,arg[4]);

  double area;

  // style of direction of interface

  if (strcmp(arg[5],"x") == 0) {
    istyle = X;
    epsilon = force->numeric(FLERR,arg[6]); 
    area = 2.0 * domain->yprd * domain->zprd;
  } else if (strcmp(arg[5],"y") == 0) {
    istyle = Y;
    epsilon = force->numeric(FLERR,arg[6]);
    area = 2.0 * domain->xprd * domain->zprd;
  } else if (strcmp(arg[5],"z") == 0) {
    istyle = Z;
    epsilon = force->numeric(FLERR,arg[6]);
    area = 2.0 * domain->yprd * domain->zprd;
  } else error->all(FLERR,"Illegal fix tasm command"); 

  if (epsilon <= 0.0) error->all(FLERR,"Illegal fix tasm command");
  if (epsilon == 1.0) error->all(FLERR,"Illegal fix tasm command");

  // set fixed-point to default = center of cell

  fixedpoint[0] = 0.5*(domain->boxlo[0]+domain->boxhi[0]);
  fixedpoint[1] = 0.5*(domain->boxlo[1]+domain->boxhi[1]);
  fixedpoint[2] = 0.5*(domain->boxlo[2]+domain->boxhi[2]);

 // optional keywords

  tailflag = 0;

  fp = NULL;

    if (strcmp(arg[7], "tail") == 0) {
      if (strcmp(arg[8],"no") == 0) tailflag = 0;
      else if (strcmp(arg[8],"yes") == 0) tailflag = 1;
      else error->all(FLERR,"Illegal optional keyword in fix tasm"); 
    }
    if (strcmp(arg[9],"file") == 0) {
      fp = fopen(arg[10],"w");
        if (comm->me == 0) {
          fprintf(fp,"# Print output for fix tasm\n");
          fprintf(fp,"# Temperature = %f\n",temp_tasm);
          fprintf(fp,"# Interfacial area = %f\n",area);     
          fprintf(fp,"# Perturbation area = %f\n",fabs(epsilon)*area);
          fprintf(fp,"# Step U_ref U_up U_down\n");
        }
    }

  x_orig = NULL;
  v_orig = NULL;
  f_orig = NULL;
  peatom_orig = NULL;
  pvatom_orig = NULL;

  allocate_storage();

  nrigid = 0;
  rfix = NULL; 
  
}

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

FixTASM::~FixTASM()
{
  delete [] vector;
  delete [] rfix;

  deallocate_storage();

  if (comm->me == 0) fclose(fp);
} 

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

int FixTASM::setmask()
{
  int mask = 0;
  mask |= END_OF_STEP;
  mask |= THERMO_ENERGY;
  return mask;
}

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

void FixTASM::init()
{
  if (domain->triclinic)
    error->all(FLERR,"Cannot use fix tasm with triclinic box");

  if (tailflag) {
    if (force->pair->tail_flag == 0)
      error->all(FLERR,"Fix tasm when pair style does not "
                 "compute tail corrections");
  }

 

  // detect if any rigid fixes exist so rigid bodies move when box is remapped
  // rfix[] = indices to each fix rigid

  delete [] rfix;
  nrigid = 0;
  rfix = NULL;

  for (int i = 0; i < modify->nfix; i++)
    if (modify->fix[i]->rigid_flag) nrigid++;
  if (nrigid) {
    rfix = new int[nrigid];
    nrigid = 0;
    for (int i = 0; i < modify->nfix; i++)
      if (modify->fix[i]->rigid_flag) rfix[nrigid++] = i;
  }
}

/* ----------------------------------------------------------------------
   free energy perturbation 
------------------------------------------------------------------------- */

void FixTASM::end_of_step()
{
  double pe_reference,pe_perturbation;
  int i;
  double lxold,lxnew,lyold,lynew,lzold,lznew;
//  if (t == 0) return;

//  if (t%nevery /= 0) return;
  //printf("entre\n");

  //if(update->ntimestep > 0 && update->ntimestep % nevery == 0) return;

  if(update->ntimestep % nevery) return;

  eflag = 1;
  vflag = 0;

  if (atom->nmax > nmax) {  // reallocate working arrays if necessary
    deallocate_storage();
    allocate_storage();
  }

// save positions, forces, energies, virial, etc

  backup_settings();

  timer->stamp();
  if (force->pair && force->pair->compute_flag) {
    force->pair->compute(eflag,vflag);
    timer->stamp(Timer::PAIR);
    //timer->stamp(TIME_PAIR);
  }

// potential energy of reference system

  pe_reference = energy_settings();

  vector[0] = pe_reference;

// increase the interfacial area

  for (i = 0; i < 3; i++) {
      oldlo[i] = lo_orig[i];
      oldhi[i] = hi_orig[i];
  }  

  iflag = 1;

  perturbation_settings(iflag);

  lxold=oldhi[0]-oldlo[0];
  lxnew=hi_cur[0]-lo_cur[0];
  lyold=oldhi[1]-oldlo[1];
  lynew=hi_cur[1]-lo_cur[1];
  lzold=oldhi[2]-oldlo[2];
  lznew=hi_cur[2]-lo_cur[2];


//  printf("+delta box old %f %f %f box new %f %f %f step %i\n",lo_cur[0],hi_cur[0],lo_cur[1],hi_cur[1],lo_cur[2],hi_cur[2],t);//lxold,lyold,lzold,lxnew,lynew,lznew); 

// decide to build the neighbor list 

  neighbor_settings();

  timer->stamp();
  if (force->pair && force->pair->compute_flag) {
    force->pair->compute(eflag,vflag);
    timer->stamp(Timer::PAIR);
    //timer->stamp(TIME_PAIR);
  }

// potential energy of perturbation system

  pe_perturbation = energy_settings();
  
  vector[1] = pe_perturbation;
  vector[3] = exp(-(pe_perturbation - pe_reference)/(force->boltz*temp_tasm));
 
// restore positions, forces, energies, virial, etc

  restore_settings();

// decide to build the neighbor list 

  neighbor_settings();

// decrease the interfacial area

  iflag = -1;

  for (i = 0; i < 3; i++) {
      oldlo[i] = lo_orig[i];
      oldhi[i] = hi_orig[i];
  }

  perturbation_settings(iflag);

  lxold=oldhi[0]-oldlo[0];
  lxnew=hi_cur[0]-lo_cur[0];
  lyold=oldhi[1]-oldlo[1];
  lynew=hi_cur[1]-lo_cur[1];
  lzold=oldhi[2]-oldlo[2];
  lznew=hi_cur[2]-lo_cur[2];

//  printf("-delta box old %f %f %f box new %f %f %f step %i\n",lo_cur[0],hi_cur[0],lo_cur[1],hi_cur[1],lo_cur[2],hi_cur[2],t);//  lxold,lyold,lzold,lxnew,lynew,lznew);     

// decide to build the neighbor list 

  neighbor_settings();

// potential energy of perturbation system

  timer->stamp();
  if (force->pair && force->pair->compute_flag) {
    force->pair->compute(eflag,vflag);
    timer->stamp(Timer::PAIR);
    //timer->stamp(TIME_PAIR);
  }

  pe_perturbation = energy_settings();

  vector[2] = pe_perturbation;
  vector[4] = exp(-(pe_perturbation - pe_reference)/(force->boltz*temp_tasm));

  restore_settings();

  if (comm->me == 0) {
    step = update->ntimestep;
    fprintf(fp,"%i %f %f %f %f %f\n",step,vector[0],vector[1],vector[2],vector[3],vector[4]);
    fflush(fp);
  }
// decide to build the neighbor list 

  neighbor_settings();

}

/* ---------------------------------------------------------------------- 
   obtain the potential energy of the current configuration
------------------------------------------------------------------------- */

double FixTASM::energy_settings()
{
  double eng;
  double eng_system;

  eng = 0.0;

  if (force->pair)
    eng += force->pair->eng_vdwl;

  MPI_Allreduce(&eng,&eng_system,1,MPI_DOUBLE,MPI_SUM,world);

  if(tailflag) { //if (force->pair && force->pair->tail_flag) {
    double volume = domain->xprd * domain->yprd * domain->zprd;
    eng_system += force->pair->etail / volume;
  }

  return eng_system;
}

/* ----------------------------------------------------------------------
   manage storage for charge, force, energy, virial arrays
------------------------------------------------------------------------- */

void FixTASM::allocate_storage()
{
  nmax = atom->nmax;
  memory->create(x_orig,nmax,3,"tasm:x_orig");  
  memory->create(v_orig,nmax,3,"tasm:v_orig");
  memory->create(f_orig,nmax,3,"tasm:f_orig");
  memory->create(peatom_orig,nmax,"tasm:peatom_orig");
  memory->create(pvatom_orig,nmax,6,"tasm:pvatom_orig");
}

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

void FixTASM::deallocate_storage()
{
  memory->destroy(x_orig);
  memory->destroy(v_orig);
  memory->destroy(f_orig);
  memory->destroy(peatom_orig);
  memory->destroy(pvatom_orig);

  v_orig = NULL;
  x_orig = NULL;
  f_orig = NULL;
  peatom_orig = NULL;
  pvatom_orig = NULL;
}

/* ----------------------------------------------------------------------
   backup the reference system 
------------------------------------------------------------------------- */

void FixTASM::backup_settings()
{
  int i,nall,natom;

  nall = atom->nlocal + atom->nghost;
  natom = atom->nlocal;

  double **x = atom->x;
  int *mask = atom->mask;
  double **f = atom->f;
  double **v = atom->v;

  
  if (force->newton)
     natom += atom->nghost;

  for (i = 0; i < 3; i++) {
      lo_orig[i] = domain->boxlo[i];
      hi_orig[i] = domain->boxhi[i];
  }
 
  for (i = 0; i < natom; i++) {
    x_orig[i][0] = x[i][0];
    x_orig[i][1] = x[i][1];
    x_orig[i][2] = x[i][2];
    v_orig[i][0] = v[i][0];
    v_orig[i][1] = v[i][1];
    v_orig[i][2] = v[i][2];
  }

  for (i = 0; i < natom; i++) {
    f_orig[i][0] = f[i][0];
    f_orig[i][1] = f[i][1];
    f_orig[i][2] = f[i][2];
  }

  eng_vdwl_orig = force->pair->eng_vdwl;

  pvirial_orig[0] = force->pair->virial[0];
  pvirial_orig[1] = force->pair->virial[1];
  pvirial_orig[2] = force->pair->virial[2];
  pvirial_orig[3] = force->pair->virial[3];
  pvirial_orig[4] = force->pair->virial[4];
  pvirial_orig[5] = force->pair->virial[5];

  if (update->eflag_atom) {
    double *peatom = force->pair->eatom;
    for (i = 0; i < natom; i++) {
      peatom_orig[i] = peatom[i];
    } 
  }
  if (update->vflag_atom) {
    double **pvatom = force->pair->vatom;
    for (i = 0; i < natom; i++) {
      pvatom_orig[i][0] = pvatom[i][0];
      pvatom_orig[i][1] = pvatom[i][1];
      pvatom_orig[i][2] = pvatom[i][2];
      pvatom_orig[i][3] = pvatom[i][3];
      pvatom_orig[i][4] = pvatom[i][4];
      pvatom_orig[i][5] = pvatom[i][5];
    }
  }
}

/* ----------------------------------------------------------------------
   apply perturbartion the reference system, changed the shape of the box 
------------------------------------------------------------------------- */
void FixTASM:: perturbation_settings(int iflag)
{
  int i;
  double **x = atom->x;
  int *mask = atom->mask;
  int nlocal = atom->nlocal;
  double *h = domain->h;
  double scale[3];
  // convert pertinent atoms and rigid bodies to lamda coords

  if (iflag == 1) {
    epsilon = fabs(epsilon);
  } else if (iflag == -1) {
    epsilon = -1.0*fabs(epsilon);
  }

  if (allremap) domain->x2lamda(nlocal);
//  else {
//    for (i = 0; i < nlocal; i++)
//      if (mask[i] & dilate_group_bit)
//        domain->x2lamda(x[i],x[i]);
//  }

  if (nrigid)
    for (i = 0; i < nrigid; i++)
      modify->fix[rfix[i]]->deform(0);
  
  // change the box shape

  if (istyle == X) {
    scale[0] = 1.0/(1.0+epsilon);
    scale[1] = scale[2] = sqrt(1.0+epsilon);   
  } else if (istyle == Y) {
    scale[1] = 1.0/(1.0+epsilon);
    scale[0] = scale[2] = sqrt(1.0+epsilon);  
  } else {
    scale[2] = 1.0/(1.0+epsilon);
    scale[0] = scale[1] = sqrt(1.0+epsilon);
  }
  
  // reset global and local box to new size/shape

  for (i = 0; i < 3; i++) {
      domain->boxlo[i] = (oldlo[i]-fixedpoint[i])*scale[i] + fixedpoint[i];
      domain->boxhi[i] = (oldhi[i]-fixedpoint[i])*scale[i] + fixedpoint[i];
      lo_cur[i]=domain->boxlo[i];
      hi_cur[i]=domain->boxhi[i];
  }

  domain->set_global_box();
  domain->set_local_box();

  // convert pertinent atoms and rigid bodies back to box coords

  if (allremap) domain->lamda2x(nlocal);
//  else {
//    for (i = 0; i < nlocal; i++)
//      if (mask[i] & dilate_group_bit)
//        domain->lamda2x(x[i],x[i]);
//  }

  if (nrigid)
    for (i = 0; i < nrigid; i++)
      modify->fix[rfix[i]]->deform(1); 
}

/* ----------------------------------------------------------------------
   restore the reference system 
------------------------------------------------------------------------- */
void FixTASM::restore_settings()
{
  int i,nall,natom;

  nall = atom->nlocal + atom->nghost;
  natom = atom->nlocal;

  double **x = atom->x;
  int *mask = atom->mask;
  double **f = atom->f;
  double **v = atom->v;  

  if (force->newton)
    natom += atom->nghost;

  for (i = 0; i < 3; i++) {
     domain->boxlo[i]=lo_orig[i];
     domain->boxhi[i]=hi_orig[i];
  }

  domain->set_global_box();
  domain->set_local_box();

  for (i = 0; i < natom; i++) {
    x[i][0] = x_orig[i][0];
    x[i][1] = x_orig[i][1];
    x[i][2] = x_orig[i][2];
    v[i][0] = v_orig[i][0];
    v[i][1] = v_orig[i][1];
    v[i][2] = v_orig[i][2];
  }

  for (i = 0; i < natom; i++) {
    f[i][0] = f_orig[i][0];
    f[i][1] = f_orig[i][1];
    f[i][2] = f_orig[i][2];
  }

  force->pair->eng_vdwl = eng_vdwl_orig;

  force->pair->virial[0] = pvirial_orig[0];
  force->pair->virial[1] = pvirial_orig[1];
  force->pair->virial[2] = pvirial_orig[2];
  force->pair->virial[3] = pvirial_orig[3];
  force->pair->virial[4] = pvirial_orig[4];
  force->pair->virial[5] = pvirial_orig[5];

  if (update->eflag_atom) {
    double *peatom = force->pair->eatom;
    for (i = 0; i < natom; i++)
      peatom[i] = peatom_orig[i];
  }
  if (update->vflag_atom) {
    double **pvatom = force->pair->vatom;
    for (i = 0; i < natom; i++) {
      pvatom[i][0] = pvatom_orig[i][0];
      pvatom[i][1] = pvatom_orig[i][1];
      pvatom[i][2] = pvatom_orig[i][2];
      pvatom[i][3] = pvatom_orig[i][3];
      pvatom[i][4] = pvatom_orig[i][4];
      pvatom[i][5] = pvatom_orig[i][5];
    }
  }
}

/* ----------------------------------------------------------------------
   decide to build the neighbor list 
------------------------------------------------------------------------- */
void FixTASM::neighbor_settings()
{
  int nflag;

  // regular communication vs neighbor list rebuild

  nflag = neighbor->decide();

  if (nflag == 0) {
     timer->stamp();
     comm->forward_comm();
     timer->stamp(Timer::COMM);
     //timer->stamp(TIME_COMM);
  } else {
    domain->pbc();
    domain->reset_box();
    comm->setup();
    if (neighbor->style) neighbor->setup_bins();
    timer->stamp();
    comm->exchange();
//    if (verlet->sortflag && update->ntimestep >= atom->nextsort) atom->sort();
    comm->borders();
    timer->stamp(Timer::COMM);
    neighbor->build();
    timer->stamp(Timer::NEIGH);
  }

// integrate->verlet->force_clear();
}

/* ----------------------------------------------------------------------
   information about the perturbation energies r
------------------------------------------------------------------------- */

double FixTASM::compute_vector(int n)
{
  return vector[n];
}