/*
  ewald.cpp
  by Hark Lee and Wei Cai, hblee@stanford.edu, caiwei@stanford.edu
  Last Modified : Mon Aug 27 11:20:51 2007

  FUNCTION  :  Ewald Summation (classical Ewald and Particle Mesh Ewald)

  To Do:
      1.  CE Virial stress still incorrect (perhaps from Self energy contribution)
      2.  CG relax dies early when both Y and Oxygen vacancy exist
      3.  PME still needs to be included

  Note:
      This file is included by md.cpp.  This file cannot be compiled by itself.
*/

/********************************************************************/
/* Ewald Summation for Coulomb Interaction */
/********************************************************************/



//-------------------------------------------------------------------
// Classical Ewald
//-------------------------------------------------------------------
void MDFrame::CE()
{
    int i;
    CE_Real();
    CE_Rec();

    _EPOT_Ewald = _EPOT_Ewald_Real + _EPOT_Ewald_Rec + _EPOT_Ewald_Self;
    _VIRIAL_Ewald = _VIRIAL_Ewald_Real + _VIRIAL_Ewald_Rec;

    for(i=0;i<_NP;i++)
    {
        _F_Ewald[i] = _F_Ewald_Real[i] + _F_Ewald_Rec[i];
        _EPOT_IND_Ewald[i] = _EPOT_IND_Ewald_Real[i] + _EPOT_IND_Ewald_Rec[i];
    }
}

//-------------------------------------------------------------------
// Compute Real Part of Ewald Summation
//-------------------------------------------------------------------
void MDFrame::CE_Real()
{
    int i,j,ipt,jpt,isp,jsp; // integers for atom counting, atom indices, atom species
    double UC,r,r2,QQ,EXP;   // U = potential | r = radius | r2 = r^2 | ri6 = r^-6
    Vector3 sij, rij, fijC; // sij : scaled coordinates | rij : coordinates | fij : forces

    DUMP(HIG"Coulomb Potential (real)"NOR);

    // start time of real part
    CE_sREAL=clock();
    
    // refresh neighbor list
    refreshneighborlist();
    
    // initialize force, potential and virials
    _EPOT_Ewald_Real = 0.0;
    _VIRIAL_Ewald_Real.clear();

    for(i=0;i<_NP;i++)
    {
        _F_Ewald_Real[i].clear();
        _EPOT_IND_Ewald_Real[i]=0.0;
    }
    
    //-------------------------------------------------------------------
    //       Calculating Short Range Potential and Ewald Real Part
    //-------------------------------------------------------------------
    for(ipt=0;ipt<_NP;ipt++) // for all particles
    {
        /* if fixed == -1, simply ignore this atom */
        if(fixed[ipt]==-1) continue;
        for(j=0;j<nn[ipt];j++) // for all neighbors
	{
            jpt=nindex[ipt][j]; // get the atom index of j-th neighbor
            /* if fixed == -1, simply ignore this atom */
            if(fixed[jpt]==-1) continue;
            
            if(ipt>=jpt) continue; // only if i<j -> avoid redundancy
            sij=_SR[jpt]-_SR[ipt]; // scaled interatomics distance
            sij.subint(); // apply periodic boundary condition (PBC)
            rij=_H*sij; // get the interatomic distance vector (with PBC)
            r2=rij.norm2(); // get the square of interatomic distance
            r=sqrt(r2); // interatomic distance
            
            if(r<=Ewald_Rcut) // if the interatomic distance is shorter than the cut-off radius
	    {
                // get the species indices
                isp=species[ipt]; // species of ipt-atom
                jsp=species[jpt]; // species of jpt-atom
                QQ=P_CLMB * _ATOMCHARGE[isp]*_ATOMCHARGE[jsp]; // get the charge term
                
                // calculate potential
                UC=(erfc(Ewald_Alpha*r))/r*QQ;
		
                // calculate force
                EXP=(2.0/M_SQRT_PI)*Ewald_Alpha*exp(-SQR(Ewald_Alpha*r))*QQ;
                fijC=rij*((UC+EXP)/r2);
                
                // add pair terms to the accumulative sum
                _F_Ewald_Real[ipt]-=fijC; _F_Ewald_Real[jpt]+=fijC;

                _EPOT_Ewald_Real+=UC;

                _EPOT_IND_Ewald_Real[ipt]+=UC*0.5;
                _EPOT_IND_Ewald_Real[jpt]+=UC*0.5;
                _VIRIAL_Ewald_Real.addnvv(1.,fijC,rij);
	    }
	}
    }
    
    // end time of real part
    CE_eREAL=clock();
}


void MDFrame::CE_Rec()
{    
    int i, ik, kal, kbl;
    int ka, kb, kc, np_real;
    double k2, kx, ky, kz;
    double pfac, tmp_H_L1, tmp_H_L2, tmp_H_L3;
    double UE, EXP;
    EComplex sincosk;

     // start time of reciprocal part
    CE_sREC=clock();
    
    // initialize force, potential and virials
   _EPOT_Ewald_Rec = 0.0;
    _VIRIAL_Ewald_Rec.clear();

    for(i=0;i<_NP;i++)
    {
        _F_Ewald_Rec[i].clear();
        _EPOT_IND_Ewald_Rec[i]=0.0;
    }
    
    // recalculate the system volume (in case the box is changing)
    if(conj_fixbox==0)
    {
        tmp_H_L1=sqrt(SQR(_H[0][0])+SQR(_H[1][0])+SQR(_H[2][0]));
        tmp_H_L2=sqrt(SQR(_H[0][1])+SQR(_H[1][1])+SQR(_H[2][1]));
        tmp_H_L3=sqrt(SQR(_H[0][2])+SQR(_H[1][2])+SQR(_H[2][2]));
        Ewald_cell_V=_H.det();        
    }
    else
    {
        tmp_H_L1 = Ewald_H_L1;
        tmp_H_L2 = Ewald_H_L2;
        tmp_H_L3 = Ewald_H_L3;
    }
    
    // precalculate exp(iKx*rx), exp(iKy*ry), exp(iKz*rz)
    CE_fillSinCosTables();
    
    kal=-1; // an impossible value
    kbl=0;  // Just to avoid a warning
    for(ik=0;ik<CE_nKV;ik++) // for all K-vectors
    {
        // kx, ky, kz calculated in enumKV()
        ka=CE_KV[ik].a;
        kb=CE_KV[ik].b;
        kc=CE_KV[ik].c;

        if(conj_fixbox==0)
        {
            kx=CE_KV[ik].kx*Ewald_H_L1/tmp_H_L1;
            ky=CE_KV[ik].ky*Ewald_H_L1/tmp_H_L2;
            kz=CE_KV[ik].kz*Ewald_H_L1/tmp_H_L3;
            k2=kx*kx+ky*ky+kz*kz;
        }
        else
        {
            kx=CE_KV[ik].kx;
            ky=CE_KV[ik].ky;
            kz=CE_KV[ik].kz;
            k2=CE_KV[ik].k2;
        }
        
        //quick sckk [ Qj*exp(iKr) ] calculation
        if(kal!=ka || kbl!=kb)  // Re-calculate sc_axby
        {
            kal=ka;
            kbl=kb;
            if(kb>=0) // if ky>=0
		for(i=0;i<_NP;i++) // for all atoms
                    EComplex::Mul(CE_scx[ka][i], CE_scy[kb][i], CE_sc_axby[i]);
            else      // if ky<0
		for(i=0;i<_NP;i++) // for all atoms
                    EComplex::MulConjugate(CE_scx[ka][i], CE_scy[-kb][i], CE_sc_axby[i]);
        }
        // initialize structure factor 
        sincosk.Re=0.0;
        sincosk.Im=0.0;
        
        if(kc>=0) // if kz>=0
	    for(i=0;i<_NP;i++) 
            {
                /* if fixed == -1, simply ignore this atom */
                if(fixed[i]==-1) continue;
                
		EComplex::Mul(CE_sc_axby[i], CE_scz[kc][i], CE_sckk[i]);
		CE_sckk[i]*=P_SQRT_CLMB*_ATOMCHARGE[species[i]];
		sincosk+=CE_sckk[i];
            }
        else      // if kz<0
	    for(i=0;i<_NP;i++)
            {
                /* if fixed == -1, simply ignore this atom */
                if(fixed[i]==-1) continue;
                
		EComplex::MulConjugate(CE_sc_axby[i], CE_scz[-kc][i], CE_sckk[i]);
		CE_sckk[i]*=P_SQRT_CLMB*_ATOMCHARGE[species[i]];
		sincosk+=CE_sckk[i];
            }

        // calculate the reciprocal energy pre-factor
        EXP=((8.0*M_PI)/Ewald_cell_V)*exp(-k2/(4.0*SQR(Ewald_Alpha)))/k2;

        // calculate the reciprocal force
        for(i=0;i<_NP;i++)
        {
            /* if fixed == -1, simply ignore this atom */
            if(fixed[i]==-1) continue;
            
            double ad=EXP*(CE_sckk[i].Im*sincosk.Re-CE_sckk[i].Re*sincosk.Im);
            _F_Ewald_Rec[i].x+=ad*((double)kx);
            _F_Ewald_Rec[i].y+=ad*((double)ky);
            _F_Ewald_Rec[i].z+=ad*((double)kz);
        }
        pfac=1.0/k2+1/(4.0*SQR(Ewald_Alpha));
        UE=EXP/2.0*sincosk.Norm2();
        
        // Potential Energy Calculation        
        _EPOT_Ewald_Rec+=UE;

        /* Temporary: distribute rec energy equally to all atoms */
        np_real = 0;
        for(i=0;i<_NP;i++)
            if(fixed[i]!=-1) np_real++;
            
        for(i=0;i<_NP;i++)
        {        
            /* if fixed == -1, simply ignore this atom */
            if(fixed[i]==-1) continue;
            
            _EPOT_IND_Ewald_Rec[i]=UE/np_real;
        }
        
        // Virial Stress Calculation
        _VIRIAL_Ewald_Rec[0][0]+=UE*(1.0-2.0*((double)(kx*kx))*pfac); //xx
        _VIRIAL_Ewald_Rec[1][1]+=UE*(1.0-2.0*((double)(ky*ky))*pfac); //yy
        _VIRIAL_Ewald_Rec[2][2]+=UE*(1.0-2.0*((double)(kz*kz))*pfac); //zz
        _VIRIAL_Ewald_Rec[0][1]+=UE*(-2.0*((double)(kx*ky))*pfac);  //xy
        _VIRIAL_Ewald_Rec[1][0]+=UE*(-2.0*((double)(kx*ky))*pfac);  //yx
        _VIRIAL_Ewald_Rec[0][2]+=UE*(-2.0*((double)(kx*kz))*pfac);  //xz
        _VIRIAL_Ewald_Rec[2][0]+=UE*(-2.0*((double)(kx*kz))*pfac);  //zx
        _VIRIAL_Ewald_Rec[1][2]+=UE*(-2.0*((double)(ky*kz))*pfac);  //yz
        _VIRIAL_Ewald_Rec[2][1]+=UE*(-2.0*((double)(ky*kz))*pfac);  //zy
    }
    //INFO("[CE] Reciprocal Energy Calculated by CE = "<<_EREC);
    
    // start time of reciprocal part
    CE_eREC=clock();
}

//-------------------------------------------------------------------
// Initialize Classical Ewald 
//-------------------------------------------------------------------
void MDFrame::CE_init()
{
    Matrix33 hinv;
    double tmp1, tmp2, tmp3, tmp4, tmp_RLIST;
    int i;
    
    // Calculate inverse matrix and length of lattice vectors
    hinv = _H.inv();
    Ewald_H_L1=sqrt(SQR(_H[0][0])+SQR(_H[1][0])+SQR(_H[2][0]));
    Ewald_H_L2=sqrt(SQR(_H[0][1])+SQR(_H[1][1])+SQR(_H[2][1]));
    Ewald_H_L3=sqrt(SQR(_H[0][2])+SQR(_H[1][2])+SQR(_H[2][2]));

    // Allocate real space force
    Realloc(_F_Ewald,     Vector3,_NP);
    Realloc(_F_Ewald_Real,Vector3,_NP);
    Realloc(_F_Ewald_Rec ,Vector3,_NP);

    Realloc(_EPOT_IND_Ewald,     double,_NP);
    Realloc(_EPOT_IND_Ewald_Real,double,_NP);
    Realloc(_EPOT_IND_Ewald_Rec ,double,_NP);
    
    bindvar("F_Ewald",_F_Ewald,DOUBLE);
    bindvar("F_Ewald_Real",_F_Ewald_Real,DOUBLE);
    bindvar("F_Ewald_Rec", _F_Ewald_Rec, DOUBLE);
    
    bindvar("EPOT_IND_Ewald",_EPOT_IND_Ewald,DOUBLE);
    bindvar("EPOT_IND_Ewald_Real",_EPOT_IND_Ewald_Real,DOUBLE);
    bindvar("EPOT_IND_Ewald_Rec", _EPOT_IND_Ewald_Rec, DOUBLE);
    
    // Allocating sc_axby and sckk
    Realloc(CE_sckk,EComplex,(_NP*2));
    CE_sc_axby=CE_sckk+_NP;

    // get the system volume
    Ewald_cell_V=_H.det();

    //-------------------------------------------------------------------    
    // Detemine the Ewald accuracy
    //-------------------------------------------------------------------
    // Ewaldprecision=sqrt(-ln(error*Ewaldprecision^2))
    //-------------------------------------------------------------------
    // sqrt(-ln(10^-4)) = 3.03485425877029
    // sqrt(-ln(10^-5)) = 3.39307021220756
    // sqrt(-ln(10^-6)) = 3.71692218884984
    // sqrt(-ln(10^-7)) = 4.01473481701573
    // sqrt(-ln(10^-8)) = 4.29193205257869
    // sqrt(-ln(10^-9)) = 4.55228138815544
    // sqrt(-ln(10^-10))= 4.79852591218808
    //-------------------------------------------------------------------

    // check if Alpha has non-zero value
    if((Ewald_option_Alpha==2)&&(Ewald_option_Alpha==0))
    {
        WARNING("Ewald Parameter Alpha is Not Specified -> Check Your Script File");
    }

    // determine Alpha (for CE only)
    if(Ewald_option_Alpha==0)
    {
        // get the optimal alpha
        Ewald_Alpha = M_SQRT_PI*cbrt(sqrt(Ewald_time_ratio*_NP)/Ewald_cell_V);
        Ewald_Rcut = Ewald_precision / Ewald_Alpha;

        if (_SKIN <= 0) _SKIN = 1.0;
        tmp_RLIST = Ewald_Rcut +_SKIN;
        tmp1=Ewald_H_L1/3.0;
        tmp2=Ewald_H_L2/3.0;
        tmp3=Ewald_H_L3/3.0;
        tmp4=tmp_RLIST;

        if(tmp4>tmp1){tmp4=tmp1;}
        if(tmp4>tmp2){tmp4=tmp2;}
        if(tmp4>tmp3){tmp4=tmp3;}
        
        if(tmp_RLIST>tmp4)
        {
            WARNING("PP_RC Computed from Optimum Alpha is Too Big");
            WARNING("Automatically Setting _RLIST Equal to 1/3 of Minimum Cell Length");
            Ewald_Rcut = (tmp4*0.95) - _SKIN;
        }
    }
    
    if(Ewald_option_Alpha==1)
    {
        // get the controlled alpha
        Ewald_Alpha = Ewald_precision / Ewald_Rcut;
    }
    
    // Update Verlet list cut-off
    _RLIST = Ewald_Rcut + _SKIN;
    
    // calculate sum of square charge
    Ewald_qSQRsum=0.0;
    for(i=0;i<_NP;i++)
    {
        /* if fixed == -1, simply ignore this atom */
        if(fixed[i]==-1) continue;
        Ewald_qSQRsum+=SQR(P_SQRT_CLMB*_ATOMCHARGE[species[i]]);
    }

    // compute the self-interaction energy
    _EPOT_Ewald_Self = - Ewald_Alpha / M_SQRT_PI * Ewald_qSQRsum;

    // determine the cutoff radius
    CE_Kc = 2*Ewald_Alpha*Ewald_precision; // reciprocal-space
      
    // Calculate the K-vectors
    CE_enumKV();

    // Print Out Ewald Information
    INFO("------------------------------------------------------------------------------");
    INFO("------------------------------------------------------------------------------");
    INFO("              Classical Ewald Parameters [Long Range Interaction] ");
    INFO("------------------------------------------------------------------------------");
    //INFO("    Real Space Cutoff Radius (Rc)                 = "<<PP_RC);
    INFO("    Ewald Fourier Part Cutoff Radius (Kc)         = "<<CE_Kc);
    INFO("    Ewald Gaussian Factor (Alpha)                 = "<<Ewald_Alpha);
    INFO("    Total Number of K-Vectors                     = "<<CE_nKV);
    INFO("    Max Number of K-Vectors                       = "<<CE_MKV);
    INFO("    Max K-Points in A-direction                   = "<<CE_Mka);
    INFO("    Max K-Points in B-direction                   = "<<CE_Mkb);
    INFO("    Max K-Points in C-direction                   = "<<CE_Mkc);
    INFO("------------------------------------------------------------------------------");
    INFO("------------------------------------------------------------------------------");
}

//-------------------------------------------------------------------
//  Clear Classical Ewald
//-------------------------------------------------------------------
void MDFrame::CE_clear()
{
    // free variables used in the Classical Ewald summation    
    if(CE_scx) { free(CE_scx[0]); free(CE_scx); }
    if(CE_sckk){ free(CE_sckk); free(CE_KV); }
}

//-------------------------------------------------------------------
// Filling Sin & Cos Table
//-------------------------------------------------------------------
// calculate exp(-iKx*rx), exp(-iKy*ry), exp(-iKz*rz) 
// for all atoms and all K-vectors, should be updated every time step
//-------------------------------------------------------------------
void MDFrame::CE_fillSinCosTables()
{
    int i,ki;
    EComplex *sc1;
    for(i=0;i<_NP;i++)
    {
        // when K-vector is zero vector (kx=0, ky=0, kz=0)
        CE_scx[0][i].Re=1.0; CE_scx[0][i].Im=0.0;
        CE_scy[0][i].Re=1.0; CE_scy[0][i].Im=0.0;
        CE_scz[0][i].Re=1.0; CE_scz[0][i].Im=0.0;
    }
    if(CE_Mka>0) 
    {
        // calculate the exp(-ikx*rx) when kx=1 (for all atoms)
        sc1=CE_scx[1];
        for(i=0;i<_NP;i++)
            sc1[i].ExpI(2.0*M_PI*_SR[i].x);
        // calculate all of the term for kx>=2 (for all atoms)
        for(ki=2;ki<=CE_Mka;ki++) for(i=0;i<_NP;i++)
            EComplex::Mul(CE_scx[ki-1][i], sc1[i], CE_scx[ki][i]);
    }
    if(CE_Mkb>0) 
    {
        // calculate the exp(-iky*ry) when ky=1 (for all atoms)
        sc1=CE_scy[1];
        for(i=0;i<_NP;i++)
	    sc1[i].ExpI(2.0*M_PI*_SR[i].y);
        // calculate all of the term for ky>=2 (for all atoms)
        for(ki=2;ki<=CE_Mkb;ki++) for(i=0;i<_NP;i++)
            EComplex::Mul(CE_scy[ki-1][i], sc1[i], CE_scy[ki][i]);
    }
    if(CE_Mkc>0) 
    {
        // calculate the exp(-ikz*rz) when kz=1 (for all atoms)
        sc1=CE_scz[1];
        for(i=0;i<_NP;i++)
	    sc1[i].ExpI(2.0*M_PI*_SR[i].z);
        // calculate all of the terms when kz=>2 (for all atoms)
        for(ki=2;ki<=CE_Mkc;ki++) for(i=0;i<_NP;i++)
            EComplex::Mul(CE_scz[ki-1][i], sc1[i], CE_scz[ki][i]);
    }
}

//-------------------------------------------------------------------
// Enumerate K-Vectors
//-------------------------------------------------------------------
void MDFrame::CE_enumKV()
{
    //------------------------------------------------------------------- 
    //  In this function, following variables are prepared
    //  Mka, Mkb, Mkc; MKV, nKV, KV
    //-------------------------------------------------------------------
    const double INHOM=1.1; //The expected error in k-point estimation
    
    EComplex *p;
    int i;
    double kx, ky, kz, k2;
    int nka, nkb, nkc, ika, ikb, ikc;
        
    // Kc = (2*pi/Li)*ni and ni = Kc*Li/(2*pi)
    // So we can multiply Kc and 1/Gi to get the integer ni
    nka=(int)(CE_Kc*Ewald_H_L1/(2.0*M_PI));
    nkb=(int)(CE_Kc*Ewald_H_L2/(2.0*M_PI));
    nkc=(int)(CE_Kc*Ewald_H_L3/(2.0*M_PI));
    
    // Get the max-number of K-vectors (half sphere)
    ika=(int)(INHOM/3.0*CE_Kc*CE_Kc*CE_Kc*Ewald_cell_V/(2.0*M_PI*2.0*M_PI));

    // Allocate K-vectors
    if(CE_MKV < ika) {CE_MKV=ika; Realloc(CE_KV,K_rec,CE_MKV);}
    
    // find K-vectors
    CE_nKV=0;
    CE_Mka=CE_Mkb=CE_Mkc=0;
    for(ika=0;ika<nka+1;ika++)
        for(ikb= (ika==0? 0:-nkb);ikb<=nkb;ikb++)
            for(ikc=((ika==0 && ikb==0)?1:-nkc);ikc<=nkc;ikc++)
            {
                //Notice that this (kx,ky,kz) is row vector, mul in left (??)
                kx=(2.0*M_PI)*(ika/Ewald_H_L1);
                ky=(2.0*M_PI)*(ikb/Ewald_H_L2);
                kz=(2.0*M_PI)*(ikc/Ewald_H_L3);
                k2=kx*kx+ky*ky+kz*kz;
                if(k2 > CE_Kc*CE_Kc) continue;
                CE_KV[CE_nKV].kx=kx;
                CE_KV[CE_nKV].ky=ky;
                CE_KV[CE_nKV].kz=kz;
                CE_KV[CE_nKV].k2=k2;
                CE_KV[CE_nKV].a=ika;
                CE_KV[CE_nKV].b=ikb;
                CE_KV[CE_nKV].c=ikc;
                CE_nKV++;
                
                if(CE_nKV >= CE_MKV) INFO("Error : enumKV() -> Out of KV space");
                if(CE_Mka < ika) CE_Mka=ika;
                if(CE_Mkb < abs(ikb)) CE_Mkb=abs(ikb); // Mkb can be negative
                if(CE_Mkc < abs(ikc)) CE_Mkc=abs(ikc); // Mkc can be negative
            }
    if(CE_scx==NULL) p=0;
    else p=CE_scx[0]; // assign p to scx[0]
    Realloc(p,EComplex,(_NP*(CE_Mka+CE_Mkb+CE_Mkc+3)));
    INFO("CE - CE_enumKV() : Number of Sincos Table Entries = "<<(_NP*(CE_Mka+CE_Mkb+CE_Mkc+3)));
    if(p==NULL) INFO("Error : CE_enumKV() -> Out of memory allocating sin/cos tables");
    Realloc(CE_scx,EComplex*,(CE_Mka+CE_Mkb+CE_Mkc+3)); // need to study more this part
    if(CE_scx==NULL) INFO("Error : CE_enumKV() -> Out of memory allocating sin/cos tables");
    CE_scy=CE_scx+(CE_Mka+1);
    CE_scz=CE_scy+(CE_Mkb+1);
    
    CE_scx[0]=p; 
    for(i=1;i<=CE_Mka;i++) CE_scx[i]=CE_scx[i-1]+_NP;
    CE_scy[0]=CE_scx[0]+(CE_Mka+1)*_NP;
    for(i=1;i<=CE_Mkb;i++) CE_scy[i]=CE_scy[i-1]+_NP;
    CE_scz[0]=CE_scy[0]+(CE_Mkb+1)*_NP;
    for(i=1;i<=CE_Mkc;i++) CE_scz[i]=CE_scz[i-1]+_NP;
}
