
/*******************************************************************************
*
* File fgcr.c
*
* Copyright (C) 2005 Martin Luescher
*
* This software is distributed under the terms of the GNU General Public
* License (GPL)
*
* Generic flexible GCR inversion program for operators Op acting on spinor
* fields
*
* The externally accessible function defined in this file is
*
*   double fgcr(int vol,int icom,void (*Op)(spinor_dble *pk,spinor_dble *pl),
*               void (*Mj)(int j,spinor *pr,spinor *pz,spinor *paz),
*               spinor *ws[],spinor_dble *wsd[],int nkv,int nmx,double res,
*               spinor_dble *pk,spinor_dble *pl,int *status)
*     Solution of the equation Op*pl[]=pk[] for given pk[]. The parameters
*     vol, nkv, nmx and res denote the number of spinors in the fields, the
*     number of Krylov vectors to be generated before a restart, the maximal
*     total number of Krylov vectors and the desired reduction of the norm
*     of the residue. Scalar products etc. are globally summed if (and only
*     if) icom=1. A workspace of 2*nkv+1 single- and 1 double-precision
*     fields is needed that are pointed to by the arrays ws[] and wsd[]. On
*     exit status reports the total number of Krylov vectors that were used,
*     or a negative value if the program failed. In all cases the norm of
*     the last residue is returned
*
* Notes:
*
* This program uses single-precision arithmetic to reduce the execution
* time but obtains the solution with double-precision accuracy
*
* The programs for the operator and the preconditioner must be externally
* defined and are assumed to have the following properties:
*
*   void Op(spinor_dble *pk,spinor_dble *pl)
*     Application of the operator Op to the double-precision field pk[]
*     and assignement of the result to pl[]. On exit the field pk[] is
*     unchanged
*
*   void Mj(int j,spinor *pr,spinor *pz,spinor *paz)
*     Application of the preconditioner M_j to the current residue pr[]
*     at step j of the algorithm. On output pr[] is unchanged, while pz[]
*     and paz[] contain the new search direction z and the field Op*z
*     respectively
*
*******************************************************************************/

#define FGCR_C

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#include "mpi.h"
#include "su3.h"
#include "start.h"
#include "linalg.h"
#include "global.h"

#define PRECISION_LIMIT ((double)(100.0f*FLT_EPSILON))

static int nsp,com,nkvm=0;
static float *be;
static complex *al,*ga,*c;
static double rn;
static spinor **psq,**psz,*psr;
static spinor_dble *pdb,*pdx,*pdw;


static int alloc_arrays(int nkv)
{
   if (nkvm>0)
   {
      afree(al);
      afree(be);
   }

   al=amalloc(nkv*(nkv+2)*sizeof(complex),3);
   be=amalloc(nkv*sizeof(float),3);

   if ((al==NULL)||(be==NULL))
      return 1;

   ga=al+nkv*nkv;
   c=ga+nkv;
   nkvm=nkv;

   return 0;
}


static void fgcr_init(int vol,int icom,int nkv,spinor *ws[],spinor_dble *wsd[],
                      spinor_dble *pk,spinor_dble *pl)
{
   nsp=vol;
   com=icom;

   psz=ws;
   psr=ws[nkv];
   psq=ws+nkv+1;

   pdb=pk;
   pdx=pl;
   pdw=wsd[0];

   set_sd2zero(nsp,pdx);
   assign_sd2s(nsp,pdb,psr);

   rn=(double)(norm_square(nsp,com,psr));
   rn=sqrt(rn);
}


static void fgcr_reset(void (*Op)(spinor_dble *pk,spinor_dble *pl))
{
   float *rsr,*rsm;
   double *rdb,*rdw;

   (*Op)(pdx,pdw);

   rdb=(double*)(pdb);
   rdw=(double*)(pdw);
   rsm=(float*)(psr+nsp);

   for (rsr=(float*)(psr);rsr<rsm;rsr+=8)
   {
      *(rsr  )=(float)((*(rdb  ))-(*(rdw  )));
      *(rsr+1)=(float)((*(rdb+1))-(*(rdw+1)));
      *(rsr+2)=(float)((*(rdb+2))-(*(rdw+2)));
      *(rsr+3)=(float)((*(rdb+3))-(*(rdw+3)));
      *(rsr+4)=(float)((*(rdb+4))-(*(rdw+4)));
      *(rsr+5)=(float)((*(rdb+5))-(*(rdw+5)));
      *(rsr+6)=(float)((*(rdb+6))-(*(rdw+6)));
      *(rsr+7)=(float)((*(rdb+7))-(*(rdw+7)));

      rdb+=8;
      rdw+=8;
   }

   rn=(double)(norm_square(nsp,com,psr));
   rn=sqrt(rn);
}


static void next_zqr(void (*Mj)(int j,spinor *pr,spinor *pz,spinor *paz),
                     int j,int nkv)
{
   int i;
   complex *z,w;
   spinor *psw;

   psw=psq[j];
   (*Mj)(j,psr,psz[j],psw);

   z=al+j;

   for (i=0;i<j;i++)
   {
      (*z)=spinor_prod(nsp,com,psq[i],psw);
      w.re=-(*z).re;
      w.im=-(*z).im;
      mulc_spinor_add(nsp,psw,psq[i],w);
      z+=nkv;
   }

   be[j]=normalize(nsp,com,psw);

   z=ga+j;
   (*z)=spinor_prod(nsp,com,psw,psr);
   w.re=-(*z).re;
   w.im=-(*z).im;
   mulc_spinor_add(nsp,psr,psw,w);

   rn=(double)(norm_square(nsp,com,psr));
   rn=sqrt(rn);
}


static void add_to_x(int nkv,int n)
{
   int k,j,kj;
   float r,*rsx,*rsm;
   complex sm;
   double *rdx;

   set_s2zero(nsp,psr);

   for (k=(n-1);k>=0;k--)
   {
      sm.re=ga[k].re;
      sm.im=ga[k].im;

      for (j=(k+1);j<n;j++)
      {
         kj=k*nkv+j;
         sm.re-=(al[kj].re*c[j].re-al[kj].im*c[j].im);
         sm.im-=(al[kj].re*c[j].im+al[kj].im*c[j].re);
      }

      r=1.0f/be[k];
      c[k].re=sm.re*r;
      c[k].im=sm.im*r;

      mulc_spinor_add(nsp,psr,psz[k],c[k]);
   }

   rdx=(double*)(pdx);
   rsm=(float*)(psr+nsp);

   for (rsx=(float*)(psr);rsx<rsm;rsx+=8)
   {
      *(rdx  )+=(double)(*(rsx  ));
      *(rdx+1)+=(double)(*(rsx+1));
      *(rdx+2)+=(double)(*(rsx+2));
      *(rdx+3)+=(double)(*(rsx+3));
      *(rdx+4)+=(double)(*(rsx+4));
      *(rdx+5)+=(double)(*(rsx+5));
      *(rdx+6)+=(double)(*(rsx+6));
      *(rdx+7)+=(double)(*(rsx+7));

      rdx+=8;
   }
}


double fgcr(int vol,int icom,void (*Op)(spinor_dble *pk,spinor_dble *pl),
            void (*Mj)(int j,spinor *pr,spinor *pz,spinor *paz),
            spinor *ws[],spinor_dble *wsd[],int nkv,int nmx,double res,
            spinor_dble *pk,spinor_dble *pl,int *status)
{
   int nv,j,iprms[3];
   double rn_old,tol,dprms[1];

   if ((icom==1)&&(NPROC>1))
   {
      iprms[0]=vol;
      iprms[1]=nkv;
      iprms[2]=nmx;
      dprms[0]=res;

      MPI_Bcast(iprms,3,MPI_INT,0,MPI_COMM_WORLD);
      MPI_Bcast(dprms,1,MPI_DOUBLE,0,MPI_COMM_WORLD);

      error((iprms[0]!=vol)||(iprms[1]!=nkv)||(iprms[2]!=nmx)||
            (dprms[0]!=res),1,"fgcr [fgcr.c]","Parameters are not global");

      error_root((vol<=0)||(nkv<1)||(nmx<1)||(res<=DBL_EPSILON),1,
                 "fgcr [fgcr.c]","Improper choice of vol,nkv,nmx or res");
      
      if (nkv>nkvm)
         error(alloc_arrays(nkv),1,"fgcr [fgcr.c]",
               "Unable to allocate auxiliary arrays");
   }
   else if (error_loc((vol<=0)||(nkv<1)||(nmx<1)||(res<=DBL_EPSILON),1,
                      "fgcr [fgcr.c]","Improper choice of vol,nkv,nmx or res"))
   {
      (*status)=0;
      return 1.0;
   }
   else if ((nkv>nkvm)&&(error_loc(alloc_arrays(nkv),1,"fgcr [fgcr.c]",
                                   "Unable to allocate auxiliary arrays")))
   {
      (*status)=0;
      return 1.0;
   }

   fgcr_init(vol,icom,nkv,ws,wsd,pk,pl);
   tol=res*rn;
   (*status)=0;

   while (rn>tol)
   {
      /* message("rn_old = %.2e\n",rn); */

      rn_old=rn;
      nv=0;

      for (j=0;j<nkv;j++)
      {
         next_zqr(Mj,j,nkv);
         nv+=1;
         (*status)+=1;
	 //printf("it %d rn = %e\n", j, rn);
	 message("j = %d, rn = %.2e it %d\n",j,rn, *status);

         if ((rn<=tol)||(rn<(PRECISION_LIMIT*rn_old))||((*status)>=nmx))
            break;
      }

      add_to_x(nkv,nv);
      fgcr_reset(Op);

      if (((*status)>=nmx)&&(rn>tol))
      {
         (*status)=-1;
         break;
      }
   }

   return rn;
}
