/* ----------------------------------------------------------------------
   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.
------------------------------------------------------------------------- */

#include "mpi.h"
#include "math.h"
#include "stdlib.h"
#include "improper_usr.h"
#include "atom.h"
#include "comm.h"
#include "neighbor.h"
#include "domain.h"
#include "force.h"
#include "update.h"
#include "math_const.h"
#include "memory.h"
#include "error.h"

using namespace LAMMPS_NS;
using namespace MathConst;

#define TOLERANCE 0.05
#define SMALL     0.001

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

ImproperUsr::ImproperUsr(LAMMPS *lmp) : Improper(lmp)
{
  writedata = 1;
}

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

ImproperUsr::~ImproperUsr()
{
  if (allocated) {
    memory->destroy(setflag);
    memory->destroy(k2);
    memory->destroy(k4);
  }
}

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

void ImproperUsr::compute(int eflag, int vflag)
{
  int i1,i2,i3,i4,n,type;
  double b12x,b12y,b12z,b13x,b13y,b13z,b14x,b14y,b14z,b134a,b134b,b134c;
  double eimproper,f1[3],f2[3],f3[3],f4[3];
  double vb1x,vb1y,vb1z,vb2x,vb2y,vb2z,vb3x,vb3y,vb3z;
  double dr1ax,dr1ay,dr1az,dr1bx,dr1by,dr1bz,dr1cx,dr1cy,dr1cz;
  double dr3ax,dr3ay,dr3az,dr3bx,dr3by,dr3bz,dr3cx,dr3cy,dr3cz;
  double dr4ax,dr4ay,dr4az,dr4bx,dr4by,dr4bz,dr4cx,dr4cy,dr4cz;
  double rsqb134,rb134,rn;
  double drsqb1341a,drsqb1341b,drsqb1341c,drsqb1343a,drsqb1343b,drsqb1343c;
  double drsqb1344a,drsqb1344b,drsqb1344c;

  eimproper = 0.0;
  if (eflag || vflag) ev_setup(eflag,vflag);
  else evflag = 0;

  double **x = atom->x;
  double **f = atom->f;
  int **improperlist = neighbor->improperlist;
  int nimproperlist = neighbor->nimproperlist;
  int nlocal = atom->nlocal;
  int newton_bond = force->newton_bond;

  for (n = 0; n < nimproperlist; n++) {
    i1 = improperlist[n][0];
    i2 = improperlist[n][1];
    i3 = improperlist[n][2];
    i4 = improperlist[n][3];
    type = improperlist[n][4];

    // geometry of 4-body

    b12x = x[i2][0] - x[i1][0];
    b12y = x[i2][1] - x[i1][1];
    b12z = x[i2][2] - x[i1][2];

    b13x = x[i3][0] - x[i1][0];
    b13y = x[i3][1] - x[i1][1];
    b13z = x[i3][2] - x[i1][2];

    b14x = x[i4][0] - x[i1][0];
    b14y = x[i4][1] - x[i1][1];
    b14z = x[i4][2] - x[i1][2];
//Virial Calculation

    vb1x = x[i1][0] - x[i2][0];
    vb1y = x[i1][1] - x[i2][1];
    vb1z = x[i1][2] - x[i2][2];

    vb2x = x[i3][0] - x[i2][0];
    vb2y = x[i3][1] - x[i2][1];
    vb2z = x[i3][2] - x[i2][2];

    vb3x = x[i4][0] - x[i3][0];
    vb3y = x[i4][1] - x[i3][1];
    vb3z = x[i4][2] - x[i3][2];
    
//Virial cal end
 
    
    //b134=b13xb14
    b134a=b13y*b14z-b13z*b14y;
    b134b=b13z*b14x-b13x*b14z;
    b134c=b13x*b14y-b13y*b14x;
    
    //Derivative of b13xb14
    dr1ax=0;
    dr1ay=-b13z+b14z;
    dr1az=-b14y+b13y;
    
    dr1bx=-b14z+b13z;
    dr1by=0;
    dr1bz=-b13x+b14x;
    
    dr1cx=-b13y+b14y;
    dr1cy=-b14x+b13x;
    dr1cz=0;
  
    dr3ax=0;
    dr3ay=-b14z;
    dr3az=b14y;
    
    dr3bx=b14z;
    dr3by=0;
    dr3bz=-b14x;
    
    dr3cx=-b14y;
    dr3cy=b14x;
    dr3cz=0;    
    
    dr4ax=0;
    dr4ay=b13z;
    dr4az=-b13y;
    
    dr4bx=-b13z;
    dr4by=0;
    dr4bz=b13x;
    
    dr4cx=b13y;
    dr4cy=-b13x;
    dr4cz=0;
        
    rsqb134 = b134a*b134a + b134b*b134b+b134c*b134c;
    if (rsqb134 < SMALL) rsqb134 = SMALL*100;
    rb134 = sqrt(rsqb134);
    rn=(b134a*b12x+b134b*b12y+b134c*b12z)/rb134;     

    //Print   
     // if (screen) {
      //      fprintf(screen,"User Style Success\n");    
     // }
    //Derivative of (b13xb14)^2
    drsqb1341a=-(b134a*0+b134b*(-b13z+b14z)+b134c*(-b14y+b13y))/pow(rsqb134,-1.5);
    drsqb1341b=-(b134a*(-b14z+b13z)+b134b*0+b134c*(-b13x+b14x))/pow(rsqb134,-1.5);
    drsqb1341c=-(b134a*(-b13y+b14y)+b134b*(-b14x+b13x)+b134c*0)/pow(rsqb134,-1.5);

    drsqb1343a=-(b134a*0+b134b*(-b14z)+b134c*(b14y))/pow(rsqb134,-1.5);
    drsqb1343b=-(b134a*(b14z)+b134b*0+b134c*(-b14x))/pow(rsqb134,-1.5);
    drsqb1343c=-(b134a*(-b14y)+b134b*(b14x)+b134c*0)/pow(rsqb134,-1.5);

    drsqb1344a=-(b134a*0+b134b*(b13z)+b134c*(-b13y))/pow(rsqb134,-1.5);
    drsqb1344b=-(b134a*(-b13z)+b134b*0+b134c*(b13x))/pow(rsqb134,-1.5);
    drsqb1344c=-(b134a*(b13y)+b134b*(-b13x)+b134c*0)/pow(rsqb134,-1.5);
    
    //Forces on each atoms
    f1[0]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr1ax*b12x+dr1ay*b12y+dr1az*b12z-b134a)/rb134+rn*(-0.5)*drsqb1341a/pow(rsqb134,2));
    f1[1]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr1bx*b12x+dr1by*b12y+dr1bz*b12z-b134b)/rb134+rn*(-0.5)*drsqb1341b/pow(rsqb134,2));
    f1[2]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr1cx*b12x+dr1cy*b12y+dr1cz*b12z-b134c)/rb134+rn*(-0.5)*drsqb1341c/pow(rsqb134,2));

    f3[0]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr3ax*b12x+dr3ay*b12y+dr3az*b12z)/rb134+rn*(-0.5)*drsqb1343a/pow(rsqb134,2));
    f3[1]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr3bx*b12x+dr3by*b12y+dr3bz*b12z)/rb134+rn*(-0.5)*drsqb1343b/pow(rsqb134,2));
    f3[2]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr3cx*b12x+dr3cy*b12y+dr3cz*b12z)/rb134+rn*(-0.5)*drsqb1343c/pow(rsqb134,2));

    f2[0]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((b134a)/rb134);
    f2[1]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((b134b)/rb134);
    f2[2]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((b134c)/rb134);

    f4[0]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr4ax*b12x+dr4ay*b12y+dr4az*b12z)/rb134+rn*(-0.5)*drsqb1344a/pow(rsqb134,2));
    f4[1]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr4bx*b12x+dr4by*b12y+dr4bz*b12z)/rb134+rn*(-0.5)*drsqb1344b/pow(rsqb134,2));
    f4[2]=-(2*k2[type]+4*k4[type]*pow(rn,2))*rn*((dr4cx*b12x+dr4cy*b12y+dr4cz*b12z)/rb134+rn*(-0.5)*drsqb1344c/pow(rsqb134,2));
       
     
     if (eflag) eimproper = k2[type]*rn*rn+k4[type]*pow(rn,4);   
   
    if (eflag && screen && n==100) {      
      fprintf(screen,"******compute:::K2: %g  K4: %g\n",k2[type],k4[type]);  
      fprintf(screen,"******compute:::energy: %g\n",eimproper);  
      fprintf(screen,"******compute:::rn:  %g rsqb134: %g  rb134: %g\n",rn,rsqb134,rb134);
      fprintf(screen,"******compute:::rn: %g  %g %g\n",f1[0],f1[1],f1[2]);
      fprintf(screen,"******compute:::f1: %g  %g %g\n",f1[0],f1[1],f1[2]); 
      fprintf(screen,"******compute:::f2: %g  %g %g\n",f2[0],f2[1],f2[2]); 
      fprintf(screen,"******compute:::f3: %g  %g %g\n",f3[0],f3[1],f3[2]); 
      fprintf(screen,"******compute:::f4: %g  %g %g\n",f4[0],f4[1],f4[2]);      
     
       }
    // apply force to each of 4 atoms

    if (newton_bond || i1 < nlocal) {
      f[i1][0] += f1[0];
      f[i1][1] += f1[1];
      f[i1][2] += f1[2];
    }

    if (newton_bond || i2 < nlocal) {
      f[i2][0] += f2[0];
      f[i2][1] += f2[1];
      f[i2][2] += f2[2];
    }

    if (newton_bond || i3 < nlocal) {
      f[i3][0] += f3[0];
      f[i3][1] += f3[1];
      f[i3][2] += f3[2];
    }

    if (newton_bond || i4 < nlocal) {
      f[i4][0] += f4[0];
      f[i4][1] += f4[1];
      f[i4][2] += f4[2];
    }

    if (evflag)
      ev_tally(i1,i2,i3,i4,nlocal,newton_bond,eimproper,f1,f3,f4,
              vb1x,vb1y,vb1z,vb2x,vb2y,vb2z,vb3x,vb3y,vb3z);
  }
}

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

void ImproperUsr::allocate()
{
  allocated = 1;
  int n = atom->nimpropertypes;
  fprintf(screen,"*************allocate:::: n: %g\n",n);
  memory->create(k2,n+1,"improper:k2");
  memory->create(k4,n+1,"improper:k4");

  memory->create(setflag,n+1,"improper:setflag");
  for (int i = 1; i <= n; i++) setflag[i] = 0;
}

/* ----------------------------------------------------------------------
   set coeffs for one type
------------------------------------------------------------------------- */

void ImproperUsr::coeff(int narg, char **arg)
{
  if (narg != 3) error->all(FLERR,"Incorrect args for improper coefficients");
  if (!allocated) allocate();

  int ilo,ihi;
  force->bounds(arg[0],atom->nimpropertypes,ilo,ihi);
  fprintf(screen,"*************coeff:::ilo: %d  ihi: %d\n",ilo,ihi);

  double k2_one = force->numeric(FLERR,arg[1]);
  double k4_one = force->numeric(FLERR,arg[2]);
  fprintf(screen,"*************coeff k2: %g  k4: %g\n",k2_one,k4_one); 
  
  // convert chi from degrees to radians

  int count = 0;
  for (int i = ilo; i <= ihi; i++) {
    fprintf(screen,"*************coeff  k2_one: %g  k4_one: %g\n",k2_one,k4_one);  
    k2[i] = k2_one;
    k4[i] = k4_one;
    fprintf(screen,"*************coeff count=%g i= %d k2: %g  k4: %g\n",count,i,k2[i],k4[i]); 
    setflag[i] = 1;
    count++;
  }

  if (count == 0) error->all(FLERR,"Incorrect args for improper coefficients");
}

/* ----------------------------------------------------------------------
   proc 0 writes out coeffs to restart file
------------------------------------------------------------------------- */

void ImproperUsr::write_restart(FILE *fp)
{
  fwrite(&k2[1],sizeof(double),atom->nimpropertypes,fp);
  fwrite(&k4[1],sizeof(double),atom->nimpropertypes,fp);
}

/* ----------------------------------------------------------------------
   proc 0 reads coeffs from restart file, bcasts them
------------------------------------------------------------------------- */

void ImproperUsr::read_restart(FILE *fp)
{
  allocate();

  if (comm->me == 0) {
    fread(&k2[1],sizeof(double),atom->nimpropertypes,fp);
    fread(&k4[1],sizeof(double),atom->nimpropertypes,fp);
  }
  MPI_Bcast(&k2[1],atom->nimpropertypes,MPI_DOUBLE,0,world);
  MPI_Bcast(&k4[1],atom->nimpropertypes,MPI_DOUBLE,0,world);

  for (int i = 1; i <= atom->nimpropertypes; i++) setflag[i] = 1;
}

/* ----------------------------------------------------------------------
   proc 0 writes to data file
------------------------------------------------------------------------- */

void ImproperUsr::write_data(FILE *fp)
{
  for (int i = 1; i <= atom->nimpropertypes; i++)
    fprintf(fp,"%d %g %g\n",i,k2[i],k4[i]);
}