/******* d_congrad5_fn_tmp.c - conjugate gradient for SU3/fermions ****/
/* MIMD version 6 */
/* Kogut-Susskind fermions -- this version for "fat plus Naik" quark
   actions.  

   This code combines d_congrad5_fn.c and d_congrad5_fn_tmp.c
   With CONGRAD_TMP_VECTORS defined, allocates temporary CG vectors in
   field-major order and uses them instead of the site-major
   temporaries.  They may be eliminated from the site structure in the
   future.

   Calls dslash_fn or dslash_fn_on_temp depending accordingly. */

/* Jim Hetrick, Kari Rummukainen, Doug Toussaint, Steven Gottlieb */
/* 10/02/01 C. DeTar Consolidated with tmp version */

/* This version looks at the initial vector every "niter" passes */
/* The source vector is in "src", and the initial guess and answer
   in "dest".  "resid" is the residual vector, and "cg_p" and "ttt" are
   working vectors for the conjugate gradient.
   niter = maximum number of iterations.
   rsqmin = desired rsq, quit when we reach rsq <= rsqmin*source_norm.
	This is different than our old definition of the stopping
	criterion.  To convert an old stopping residual to the new
	one, multiply the old one by sqrt( (2/3)/(8+2*m) )
        This is because the source is obtained from
        a random vector with average squared magnitude 3 on each site.
        Then, on 1/2 the sites, we gather and sum the eight neighboring
        random vectors and add 2*m times the local vector.
            source = M_adjoint*R, on even sites
   reinitialize after niters iterations and try once more.
   parity=EVEN = do only even sites, parity=ODD = do odd sites,
   parity=EVENANDODD = do all sites
*/
#include "generic_ks_includes.h"	/* definitions files and prototypes */
#include "../include/prefetch.h"
#define FETCH_UP 1

void cleanup_gathers(msg_tag *t1[16],msg_tag *t2[16]); /* dslash_fn_tmp.c */

#define LOOPEND
#include "../include/loopend.h"

#define QDP_Precision 'F'
#define QDP_Nc 3
#include <qdp.h>

#define dclock dclock_cpu

static int qdp_initialized=0;
static QDP_ColorMatrix *fatlinks[4], *longlinks[4];
static QDP_ColorVector *src, *dest, *ttt, *tttt, *resid, *cg_p;
static QDP_ColorVector *temp_up1[4], *temp_up2[4];
static QDP_ColorVector *temp_up31[4], *temp_up32[4];
static QDP_ColorVector *temp_dn1[4], *temp_dn2[4];
static QDP_ColorVector *temp_dn31[4], *temp_dn32[4];
static QDP_ColorVector *tempvec[4], *tempvec3[4];
static QDP_ColorVector **temp_up, **temp_dn;
static QDP_ColorVector **temp_up3, **temp_dn3;
static QDP_Shift neighbor3[4];

#define set_qdp_from_field(dest, src, TYPE, T) \
{ \
  int _i; \
  site *_s; \
  QLA_##TYPE *_temp; \
  _temp = QDP_expose_##T (dest); \
  FORALLSITES(_i,_s) { \
    memcpy((void *)&_temp[_i], F_PT(_s,src), sizeof(QLA_##TYPE)); \
  } \
  QDP_reset_##T (dest); \
}

#define set_field_from_qdp(dest, src, TYPE, T) \
{ \
  int _i; \
  site *_s; \
  QLA_##TYPE *_temp; \
  _temp = QDP_expose_##T (src); \
  FORALLSITES(_i,_s) { \
    memcpy(F_PT(_s,dest), (void *)&_temp[_i], sizeof(QLA_##TYPE)); \
  } \
  QDP_reset_##T (src); \
}

#define set4_qdp_from_field(dest, src, TYPE, T) \
{ \
  int _i, _dir; \
  site *_s; \
  QLA_##TYPE *_temp; \
  for(_dir=0; _dir<4; _dir++) { \
    _temp = QDP_expose_##T (dest[_dir]); \
    FORALLSITES(_i,_s) { \
      memcpy((void *)&_temp[_i], (void *)&src[4*_i+_dir], sizeof(QLA_##TYPE)); \
    } \
    QDP_reset_##T (dest[_dir]); \
  } \
}

#define CLEANUP \
  set_field_from_qdp(f_dest, dest, ColorVector, V);

void dslash_fn_qdp( QDP_ColorVector *src, QDP_ColorVector *dest, QDP_Subset parity );

void
init_qdp(void)
{
  int i;
  int disp[4]={0,0,0,0};

  QDP_initialize(0, NULL);
  QDP_create_layout();

  for(i=0; i<4; i++) {
    if(i!=0) disp[i-1] = 0;
    disp[i] = 3;
    neighbor3[i] = QDP_create_shift(disp);
  }

  src = QDP_create_V();
  dest = QDP_create_V();
  ttt = QDP_create_V();
  tttt = QDP_create_V();
  resid = QDP_create_V();
  cg_p = QDP_create_V();
  for(i=0; i<4; i++) {
    fatlinks[i] = QDP_create_M();
    longlinks[i] = QDP_create_M();
    temp_up1[i] = QDP_create_V();
    temp_up2[i] = QDP_create_V();
    temp_dn1[i] = QDP_create_V();
    temp_dn2[i] = QDP_create_V();
    tempvec[i] = QDP_create_V();
    temp_up31[i] = QDP_create_V();
    temp_up32[i] = QDP_create_V();
    temp_dn31[i] = QDP_create_V();
    temp_dn32[i] = QDP_create_V();
    tempvec3[i] = QDP_create_V();
  }
  qdp_initialized = 1;
}

int
ks_congrad( field_offset f_src, field_offset f_dest, float mass,
	    int niter, float rsqmin, int parity, float *final_rsq_ptr )
{
  int i;
  int iteration;	 /* counter for iterations */
  float a,b;		 /* Sugar's a,b,resid**2,last resid*2 */
  double rsq,oldrsq,pkp; /* pkp = cg_p.K.cg_p */
  float msq_x4;	         /* 4*mass*mass */
  double source_norm;	 /* squared magnitude of source vector */
  double rsqstop;	 /* stopping residual normalized by source norm */
  int l_parity;	         /* parity we are currently doing */
  int l_otherparity;	 /* the other parity */

  QDP_Subset q_parity, q_otherparity;
  QLA_Complex temp_csum;
  QLA_Real temp_sum;

/* Timing */

#ifdef CGTIME
  double dtimed,dtimec;
#endif
  double nflop;

/* debug */
#ifdef CGTIME
  dtimec = -dclock(); 
#endif

  nflop = 1187;
  if(parity==EVENANDODD) nflop *= 2;

  if(!qdp_initialized) init_qdp();

  /* if we want both parities, we will do even first. */
  switch(parity) {
  case(EVEN): l_parity=EVEN; l_otherparity=ODD; break;
  case(ODD):  l_parity=ODD; l_otherparity=EVEN; break;
  case(EVENANDODD):  l_parity=EVEN; l_otherparity=ODD; break;
  }
  if(l_parity==EVEN) {
    q_parity = QDP_even;
    q_otherparity = QDP_odd;
  } else {
    q_parity = QDP_odd;
    q_otherparity = QDP_even;
  }
  msq_x4 = 4.0*mass*mass;
  iteration = 0;

  if (!valid_fatlinks) load_fatlinks();
  if (!valid_longlinks) load_longlinks();

  temp_up = temp_up1;
  temp_dn = temp_dn1;
  temp_up3 = temp_up31;
  temp_dn3 = temp_dn31;
  set_qdp_from_field(src, f_src, ColorVector, V);
  set_qdp_from_field(dest, f_dest, ColorVector, V);
#ifdef DSLASH_TMP_LINKS
  set4_qdp_from_field(fatlinks, t_fatlink, ColorMatrix, M);
  set4_qdp_from_field(longlinks, t_longlink, ColorMatrix, M);
#else
  for(i=0; i<4; i++) {
    set_qdp_from_field(fatlinks[i], F_OFFSET(fatlink[i]), ColorMatrix, M);
    set_qdp_from_field(longlinks[i], F_OFFSET(longlink[i]), ColorMatrix, M);
  }
#endif

#ifdef CGTIME
  dtimec = -dclock(); 
#endif

  /* initialization process */
start:
  /**node0_printf("ks_congrad4: start, parity = %d\n",parity);**/
  /* ttt <-  (-1)*M_adjoint*M*dest
     resid,cg_p <- src + ttt
     rsq = |resid|^2
     source_norm = |src|^2
  */
/**if(this_node==0)if(iteration>1)printf("CONGRAD: restart rsq = %.10e\n",rsq);**/

  temp_up = temp_up1;
  temp_dn = temp_dn1;
  temp_up3 = temp_up31;
  temp_dn3 = temp_dn31;
  //fprintf(stderr,"ds11\n");
  dslash_fn_qdp(dest, tttt, q_otherparity);
  temp_up = temp_up2;
  temp_dn = temp_dn2;
  temp_up3 = temp_up32;
  temp_dn3 = temp_dn32;
  //fprintf(stderr,"ds12\n");
  dslash_fn_qdp(tttt, ttt, q_parity);

  /* ttt  <- ttt - msq_x4*src	(msq = mass squared) */
  QDP_V_meq_r_times_V(ttt, (QLA_Real *)&msq_x4, dest, q_parity);
  QDP_V_eq_V_plus_V(resid, src, ttt, q_parity);
  QDP_V_eq_V(cg_p, resid, q_parity);
  QDP_r_eq_norm2_V(&temp_sum, src, q_parity);
  source_norm = temp_sum;
  QDP_r_eq_norm2_V(&temp_sum, resid, q_parity);
  rsq = temp_sum;

  /**if(this_node==0)printf("CONGRAD: start rsq = %.10e\n",rsq);**/
  iteration++ ;  /* iteration counts number of multiplications
		    by M_adjoint*M */
  total_iters++;
  rsqstop = rsqmin * source_norm;
  /**node0_printf("congrad: source_norm = %e\n", (double)source_norm);**/
  if( rsq <= rsqstop ){
    /* if parity==EVENANDODD, set up to do odd sites and go back */
    if(parity == EVENANDODD) {
      l_parity=ODD; l_otherparity=EVEN;
      q_parity=QDP_odd; q_otherparity=QDP_even;
      parity=EVEN;	/* so we won't loop endlessly */
      iteration = 0;
      /**node0_printf("instant goto start\n"); **/
      goto start;
    }
    *final_rsq_ptr=(float)rsq;
    /**node0_printf("instant return\n"); fflush(stdout);**/
    CLEANUP;
    return (iteration);
  }
  /**pkp=0.0;
     if(mynode()==0){printf("iter=%d, rsq= %e, pkp=%e\n",
     iteration,(double)rsq,(double)pkp);fflush(stdout);}**/

  /* main loop - do until convergence or time to restart */
  /*
    oldrsq <- rsq
    ttt <- (-1)*M_adjoint*M*cg_p
    pkp <- (-1)*cg_p.M_adjoint*M.cg_p
    a <- -rsq/pkp
    dest <- dest + a*cg_p
    resid <- resid + a*ttt
    rsq <- |resid|^2
    b <- rsq/oldrsq
    cg_p <- resid + b*cg_p
  */
  do{
    oldrsq = rsq;
    pkp = 0.0;
    /* sum of neighbors */

    temp_up = temp_up1;
    temp_dn = temp_dn1;
    temp_up3 = temp_up31;
    temp_dn3 = temp_dn31;
    //fprintf(stderr,"ds21\n");
    dslash_fn_qdp( cg_p, tttt, q_otherparity );
    temp_up = temp_up2;
    temp_dn = temp_dn2;
    temp_up3 = temp_up32;
    temp_dn3 = temp_dn32;
    //fprintf(stderr,"ds22\n");
    dslash_fn_qdp( tttt, ttt, q_parity );

    /* finish computation of M_adjoint*m*p and p*M_adjoint*m*Kp */
    /* ttt  <- ttt - msq_x4*cg_p	(msq = mass squared) */
    /* pkp  <- cg_p.(ttt - msq*cg_p) */
    QDP_V_meq_r_times_V(ttt, (QLA_Real *)&msq_x4, cg_p, q_parity);
    QDP_c_eq_V_dot_V(&temp_csum, cg_p, ttt, q_parity);
    pkp = QLA_real(temp_csum);

    iteration++;
    total_iters++;

    a = (float) (-rsq/pkp);

    /* dest <- dest - a*cg_p */
    /* resid <- resid - a*ttt */
    QDP_V_peq_r_times_V(dest, (QLA_Real *)&a, cg_p, q_parity);
    QDP_V_peq_r_times_V(resid, (QLA_Real *)&a, ttt, q_parity);
    QDP_r_eq_norm2_V(&temp_sum, resid, q_parity);
    rsq = temp_sum;

    /**if(mynode()==0){printf("iter=%d, rsq= %e, pkp=%e\n",
       iteration,(double)rsq,(double)pkp);fflush(stdout);}**/

    if( rsq <= rsqstop ){
      /* if parity==EVENANDODD, set up to do odd sites and go back */
      if(parity == EVENANDODD) {
	l_parity=ODD; l_otherparity=EVEN;
	q_parity=QDP_odd; q_otherparity=QDP_even;
	parity=EVEN;	/* so we won't loop endlessly */
	iteration = 0;
	/**node0_printf("normal goto start\n"); **/
	goto start;
      }
      *final_rsq_ptr=(float)rsq;
      /**node0_printf("normal return\n"); fflush(stdout);**/
#ifdef CGTIME
      dtimec += dclock();
      if(this_node==0){printf("CONGRAD5: time = %e iters = %d mflops = %e\n",
			      dtimec,iteration,(double)(nflop*volume*iteration/(1.0e6*dtimec*numnodes())) );
      fflush(stdout);}
#endif
      CLEANUP;
      return (iteration);
    }

    b = (float)rsq/oldrsq;
    /* cg_p  <- resid + b*cg_p */
    QDP_V_eq_r_times_V_plus_V(cg_p, (QLA_Real *)&b, cg_p, resid, q_parity);

  } while( iteration%niter != 0);

  if( iteration < 5*niter ){
    /**node0_printf("try again goto start\n");**/
    goto start;
  }

  /* if parity==EVENANDODD, set up to do odd sites and go back */
  if(parity == EVENANDODD) {
    l_parity=ODD; l_otherparity=EVEN;
    q_parity=QDP_odd; q_otherparity=QDP_even;
    parity=EVEN;	/* so we won't loop endlessly */
    iteration = 0;
    goto start;
  }

  *final_rsq_ptr=rsq;
  node0_printf("CG not converged after %d iterations, res. = %e wanted %e\n",
	       iteration,rsq,rsqstop);
  fflush(stdout);
  CLEANUP;
  return(iteration);
}

/* clear an su3_vector in the lattice */
void clear_latvec(field_offset v,int parity){
register int i,j;
register site *s;
register su3_vector *vv;
    switch(parity){
	case EVEN: FOREVENSITES(i,s){
		vv = (su3_vector *)F_PT(s,v);
		for(j=0;j<3;j++){ vv->c[j].real = vv->c[j].imag = 0.0; }
	    } break;
	case ODD: FORODDSITES(i,s){
		vv = (su3_vector *)F_PT(s,v);
		for(j=0;j<3;j++){ vv->c[j].real = vv->c[j].imag = 0.0; }
	    } break;
	case EVENANDODD: FORALLSITES(i,s){
		vv = (su3_vector *)F_PT(s,v);
		for(j=0;j<3;j++){ vv->c[j].real = vv->c[j].imag = 0.0; }
	    } break;
    } 
}

/* copy an su3_vector in the lattice */
void copy_latvec(field_offset src,field_offset dest,int parity){
register int i;
register site *s;
register su3_vector *spt,*dpt;
    switch(parity){
	case EVEN: FOREVENSITES(i,s){
		s = &(lattice[i]);
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		*dpt = *spt;
	    } break;
	case ODD: FORODDSITES(i,s){
		s = &(lattice[i]);
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		*dpt = *spt;
	    } break;
	case EVENANDODD: FORALLSITES(i,s){
		s = &(lattice[i]);
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		*dpt = *spt;
	    } break;
    } 
}

/* scalar multiply and add an SU3 vector in the lattice */
void scalar_mult_add_latvec(field_offset src1,field_offset src2,
			    float scalar,field_offset dest,int parity)
{
register int i;
register site *s;
register su3_vector *spt1,*spt2,*dpt;
        FORSOMEPARITY(i,s,parity){
               spt1 = (su3_vector *)F_PT(s,src1);
                spt2 = (su3_vector *)F_PT(s,src2);
                dpt = (su3_vector *)F_PT(s,dest);
		if( i < loopend-FETCH_UP ){
		  prefetch_VVV( (su3_vector *)F_PT((s+FETCH_UP),src1),
				(su3_vector *)F_PT((s+FETCH_UP),src2),
				(su3_vector *)F_PT((s+FETCH_UP),dest) );
		}
                scalar_mult_add_su3_vector( spt1 , spt2 , scalar , dpt);
       } END_LOOP
}

void scalar2_mult_add_su3_vector(su3_vector *a, float s1, su3_vector *b, 
				 float s2, su3_vector *c){
register int i;
    for(i=0;i<3;i++){
        c->c[i].real = s1*a->c[i].real + s2*b->c[i].real;
        c->c[i].imag = s1*a->c[i].imag + s2*b->c[i].imag;
    }
}

/* scalar multiply two SU3 vector and add through the lattice */
void scalar2_mult_add_latvec(field_offset src1,float scalar1,
			     field_offset src2,float scalar2,
			     field_offset dest,int parity)
{
register int i;
register site *s;
register su3_vector *spt1,*spt2,*dpt;
        FORSOMEPARITY(i,s,parity){
		spt1 = (su3_vector *)F_PT(s,src1);
		spt2 = (su3_vector *)F_PT(s,src2);
		dpt  = (su3_vector *)F_PT(s,dest);
		if( i < loopend-FETCH_UP ){
		  prefetch_VVV((su3_vector *)F_PT((s+FETCH_UP),src1),
			       (su3_vector *)F_PT((s+FETCH_UP),src2),
			       (su3_vector *)F_PT((s+FETCH_UP),dest) );
		}
		scalar2_mult_add_su3_vector( spt1, scalar1, spt2, scalar2, dpt);
       } END_LOOP
}

/* scalar multiply an SU3 vector in the lattice */
void scalar_mult_latvec(field_offset src,float scalar,
			field_offset dest,int parity)
{
register int i;
register site *s;
register su3_vector *spt,*dpt;
    switch(parity){
	case EVEN: FOREVENSITES(i,s){
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		scalar_mult_su3_vector( spt , scalar , dpt );
	    } break;
	case ODD: FORODDSITES(i,s){
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		scalar_mult_su3_vector( spt , scalar , dpt );
	    } break;
	case EVENANDODD: FORALLSITES(i,s){
		spt = (su3_vector *)F_PT(s,src);
		dpt = (su3_vector *)F_PT(s,dest);
		scalar_mult_su3_vector( spt , scalar , dpt );
	    } break;
    } 
}

/* D_slash routine - sets dest. on each site equal to sum of
   sources parallel transported to site, with minus sign for transport
   from negative directions */
void
dslash_fn_qdp( QDP_ColorVector *src, QDP_ColorVector *dest, QDP_Subset parity )
{
  int i;
  QDP_Subset otherparity;
  QLA_Real sum;

  if(parity==QDP_even) {
    otherparity = QDP_odd;
  } else if(parity==QDP_odd) {
    otherparity = QDP_even;
  } else {
    otherparity = QDP_all;
  }

  //QDP_r_eq_norm2_V(&sum, src, otherparity);
  //fprintf(stderr,"s %g\n",sum);

#ifdef DSLASHTIMES
  dtime0 -= dclock(); 
#endif

  /* Start gathers from positive directions */
  //fprintf(stderr,"test1\n");
  for(i=0; i<4; i++) {
    QDP_V_eq_sV(temp_up[i], src, QDP_neighbor[i], QDP_forward, parity);
    QDP_V_eq_sV(temp_up3[i], src, neighbor3[i], QDP_forward, parity);
  }
  //fprintf(stderr,"test2\n");

#ifdef DSLASHTIMES
  dtime0 += dclock();
  dtime1 -= dclock(); 
#endif
#ifdef M4TIME
  dtimem -= dclock();
#endif

  /* Multiply by adjoint matrix at other sites */
  for(i=0; i<4; i++) {
    QDP_V_eq_Ma_times_V(tempvec[i], fatlinks[i], src, otherparity);
    //QDP_r_eq_norm2_V(&sum, tempvec[i], otherparity);
    //fprintf(stderr,"%g ",sum);
    QDP_V_eq_Ma_times_V(tempvec3[i], longlinks[i], src, otherparity);
    //QDP_r_eq_norm2_V(&sum, tempvec3[i], otherparity);
    //fprintf(stderr,"%g ",sum);
  }

#ifdef M4TIME
  dtimem += dclock();
  if(otherparity==QDP_all) dtimem_iters +=2; else dtimem_iters++;
#endif
#ifdef DSLASHTIMES
  dtime1 += dclock();
  dtime2 -= dclock(); 
#endif

  //fprintf(stderr,"test3\n");
  /* Start gathers from negative directions */
  for(i=0; i<4; i++) {
    QDP_V_eq_sV(temp_dn[i], tempvec[i], QDP_neighbor[i], QDP_backward, parity);
    QDP_V_eq_sV(temp_dn3[i], tempvec3[i], neighbor3[i], QDP_backward, parity);
  }
  //fprintf(stderr,"test4\n");

#ifdef DSLASHTIMES
  dtime2 += dclock();
  dtime3 -= dclock(); 
#endif
#ifdef DSLASHTIMES
  dtime3 += dclock();
  dtime4 -= dclock(); 
#endif

  /* Wait gathers from positive directions, multiply by matrix and
     accumulate */
  QDP_V_eq_M_times_V(dest, fatlinks[0], temp_up[0], parity);
  QDP_V_peq_M_times_V(dest, longlinks[0], temp_up3[0], parity);
  QDP_discard_V(temp_up[0]);
  QDP_discard_V(temp_up3[0]);
  for(i=1; i<4; i++) {
    QDP_V_peq_M_times_V(dest, fatlinks[i], temp_up[i], parity);
    QDP_V_peq_M_times_V(dest, longlinks[i], temp_up3[i], parity);
    QDP_discard_V(temp_up[i]);
    QDP_discard_V(temp_up3[i]);
  }
  //fprintf(stderr,"test5\n");

#ifdef DSLASHTIMES
  dtime4 += dclock();
  dtime5 -= dclock(); 
#endif
#ifdef DSLASHTIMES
  dtime5 += dclock();
  dtime6 -= dclock(); 
#endif

  /*------------------------------------------------------------*/
  /* Wait gathers from negative directions, accumulate (negative) */

  for(i=0; i<4; i++) {
    QDP_V_meq_V(dest, temp_dn[i], parity);
    QDP_V_meq_V(dest, temp_dn3[i], parity);
    QDP_discard_V(temp_dn[i]);
    QDP_discard_V(temp_dn3[i]);
  }
#ifdef DSLASHTIMES
  dtime6 +=dclock();
#endif
  //QDP_r_eq_norm2_V(&sum, dest, parity);
  //fprintf(stderr,"\n%g\n",sum);
  //fprintf(stderr,"test6\n");
}