/************************** gaugefix2.c *******************************/
/* Fix Coulomb or Lorentz gauge by doing successive SU(2) gauge hits */
/* Uses double precision global sums */
/* This version does automatic reunitarization at preset intervals */
/* MIMD version 6 */
/* C. DeTar 10-22-90 */
/* T. DeGrand 1993 */
/* U.M. Heller 8-31-95 */
/* C. DeTar 10-11-97 converted to generic */
/* C. DeTar 12-26-97 added automatic reunitarization */
/* C. DeTar 11-24-98 remove superfluous references to p2 (was for ks phases) */

/* Heavily modified by Kari Rummukainen 2005-6 */

/* Prototype...

void gaugefix(int gauge_dir,double relax_boost,int max_gauge_iter,
              double gauge_fix_tol, suN_matrix gauge );

   if gauge == NULL do not return the gauge

   -------------------------------------------------------------------

   NOTE: For staggered fermion applications, it is necessary to remove
   the KS phases from the gauge links before calling this procedure.
   See "rephase" in setup.c.

   -------------------------------------------------------------------
   EXAMPLE:  Fixing only the link matrices to Coulomb gauge with scratch
     space in mp (suN_matrix) and chi (suN_vector):

   gaugefix(TUP,1.5,500,1.0e-7,NULL);

   -------------------------------------------------------------------
   EXAMPLE:  Fixing Coulomb gauge with respect to the y direction
      in the staggered fermion scheme and simultaneously transforming
      the pseudofermion fields and gauge-momenta involved in updating:

   int nvector = 3;
   suN_vector * vec[3] = {g_rand, phi, xxx };
   int vector_parity[3] = { EVENODDODD, EVEN, EVEN };
   int nantiherm = 4;
   int antiherm_parity[4] = { EVENODD, EVENODD, EVENODD, EVENODD }

   rephase( OFF );
   gauge = new_latfield(suN_matrix);
   gaugefix( YUP, 1.8, 500, 2.0e-6, gauge );
   vec_fix_gauge( gauge, g_rand, EVENODD );
   vec_fix_gauge( gauge, phi, EVEN );
   vec_fix_gauge( gauge, xxx, EVEN );
   foralldir(d) ahmat_fix_gauge( gauge, mom[d], EVENODD );
   free_latfield( gauge );

   rephase( ON );

   -------------------------------------------------------------------

   gauge_dir     specifies the direction of the "time"-like hyperplane
                 for the purposes of defining Coulomb or Lorentz gauge 
      TUP    for evaluating propagators in the time-like direction
      ZUP    for screening lengths.
      -1     for Lorentz gauge
   relax_boost	   Overrelaxation parameter 
   max_gauge_iter  Maximum number of iterations 
   gauge_fix_tol   Stop if change is less than this 
*/

#include "lattice.h"
#define REUNIT_INTERVAL 20

#ifdef SU2
DOES NOT WORK YET FOR SU2
#endif

/*    CDIF(a,b)         a -= b						      */
								/*  a -= b    */
#define CDIF(a,b) { (a).real -= (b).real; (a).imag -= (b).imag; }

/* Scratch space */

void accum_gauge_hit(int i, int gauge_dir, 
		     su3_matrix *diffmat, su3_vector *sumvec )
{

/* Accumulates sums and differences of link matrices for determining optimum */
/* hit for gauge fixing */
/* Differences are kept in diffmat and the diagonal elements of the sums */
/* in sumvec  */

  register int j;
  register su3_matrix *m1;
  register int dir;

  /* Clear sumvec and diffmat */

  clear_su3mat( diffmat );
  clearvec( sumvec );

  /* Subtract upward link contributions */
    
  foralldir(dir) if (dir != gauge_dir) {
    int odir = opp_dir(dir);

    m1 = &(U[dir][i]);
    sub_su3_matrix( diffmat, m1, diffmat);
    /* Sum diagonal part */
    for(j=0; j<Ncol; j++) CSUM( sumvec->c[j], m1->e[j][j] );
  

    /* Add downward link contributions */

    m1 = &U[dir][nb(odir,i)];
    add_su3_matrix( diffmat, m1, diffmat );
    for(j=0; j<Ncol; j++) CSUM( sumvec->c[j], m1->e[j][j] );
  }
} /* accum_gauge_hit */


void do_hit(int gauge_dir, int parity, double relax_boost, su3_matrix *gauge )
{
  /* Do optimum SU(2) gauge hit for p, q subspace */

  double a0,a1,a2,a3,asq,a0sq,x,r,xdr;
  int dir,i,p,q;
  su2_matrix u;
  su3_matrix diffmat;
  su3_vector sumvec;

  /* Accumulate sums for determining optimum gauge hit - 
   * U's must have been fetched from down! */

  /* accum_gauge_hit( gauge_dir, parity, diffmat, sumvec); */

  forparity(i,parity) for (p=0; p<Ncol-1; p++) for (q=p+1; q<Ncol; q++) {

    accum_gauge_hit( i, gauge_dir, &diffmat, &sumvec );

    /* The SU(2) hit matrix is represented as a0 + i * Sum j (sigma j * aj)*/
    /* The locally optimum unnormalized components a0, aj are determined */
    /* from the current link in direction dir and the link downlink */
    /* in the same direction on the neighbor in the direction opposite dir */
    /* The expression is */
    /* a0 = Sum dir Tr Re 1       * (downlink dir + link dir) */
    /* aj = Sum dir Tr Im sigma j * (downlink dir - link dir)  j = 1,2, 3 */
    /*   where 1, sigma j are unit and Pauli matrices on the p,q subspace */
    
    a0 =  sumvec.c[p].real + sumvec.c[q].real;	
    a1 =  diffmat.e[q][p].imag + diffmat.e[p][q].imag;
    a2 = -diffmat.e[q][p].real + diffmat.e[p][q].real;
    a3 =  diffmat.e[p][p].imag - diffmat.e[q][q].imag;

    /* Over-relaxation boost */

    /* This algorithm is designed to give little change for large |a| */
    /* and to scale up the gauge transformation by a factor of relax_boost*/
    /* for small |a| */

    asq = a1*a1 + a2*a2 + a3*a3;
    a0sq = a0*a0;
    x = (relax_boost*a0sq + asq)/(a0sq + asq);
    r = sqrt((double)(a0sq + x*x*asq));
    xdr = x/r;
    /* Normalize and boost */
    a0 = a0/r; a1 = a1*xdr; a2 = a2*xdr; a3 = a3*xdr;

    /* Elements of SU(2) matrix */

    u.e[0][0] = cmplx( a0, a3);
    u.e[0][1] = cmplx( a2, a1);
    u.e[1][0] = cmplx(-a2, a1);
    u.e[1][1] = cmplx( a0,-a3);
      
    /* Do SU(2) hit on all upward links */

    foralldir(dir)
      left_su2_hit_n( &u, p, q, &U[dir][i] );
      
    /* Do SU(2) hit on all downward links */
      
    foralldir(dir)
      right_su2_hit_a( &u, p, q, &U[dir][nb(opp_dir(dir),i)] );
    
    /* accumulate gauge transform */

    if (gauge != NULL) left_su2_hit_n( &u, p, q, &gauge[i] );
  
  }
  /* Exit with modified downward links left in communications buffer */
} /* do_hit */

double get_gauge_fix_action(int gauge_dir,int parity)
{
  /* Adds up the gauge fixing action for sites of given parity */
  /* Returns average over these sites */
  /* The average is normalized to a maximum of 1 when all */
  /* links are unit matrices */

  register int dir,i,ndir;
  double gauge_fix_action;
  complex trace;

  gauge_fix_action = 0.0;
  
  /* NOTE - do not care about parity, does not matter */
  foralldir(dir) if (dir != gauge_dir) {
    forallsites(i) {      
      trace = trace_su3( &U[dir][i] );
      gauge_fix_action += (double)trace.real;
    }
  }

  /* Count number of terms to average */
  ndir = 0; foralldir(dir) if (dir != gauge_dir) ndir++;
  
  /* Sum over all sites of this parity - scatter */
  g_doublesum( &gauge_fix_action, 1 );
  
  /* Average is normalized to max 1 */
  return(gauge_fix_action /((double)(6*ndir*lattice.volume)));
} /* get_gauge_fix_action */


double gaugefixstep(int gauge_dir,double relax_boost, su3_matrix *gauge )
{
  /* Carry out one iteration in the gauge-fixing process */

  int parity;
  msg_tag *mtag[NDIM];
  register int dir;

  /* Alternate parity to prevent interactions during gauge transformation */
  g_sync();
  fflush(stdout);
  
  forbothparities(parity) {
    /* Start gathers of downward links, need for do_hit */
    foralldir(dir) {
      mtag[dir] = start_get( U[dir], opp_dir(dir), parity );
    }
      
    /* Total gauge fixing action for sites of this parity: Before 
     * this needed only for debugging */
    /* gauge_fix_action = get_gauge_fix_action( gauge_dir, parity);
     */
    foralldir(dir) wait_get(mtag[dir]);

    /* Do optimum gauge hit on various subspaces */

    do_hit( gauge_dir, parity, relax_boost, gauge );
    
    /* Scatter downward link matrices 
     * -- note, U's are left in the nb-buffers! 
     */

    foralldir(dir) mtag[dir] = start_put( U[dir], opp_dir(dir), parity );
    foralldir(dir) wait_put(mtag[dir]);

  }
  /* Total gauge fixing action for sites of this parity: After */
  /* Parity does not matter here */
  return ( get_gauge_fix_action(gauge_dir,EVENODD) );

} /* gaugefixstep */


void gaugefix(int gauge_dir,double relax_boost,int max_gauge_iter,
	      double gauge_fix_tol, su3_matrix *gauge )
{
  int gauge_iter,i;
  double current_av, old_av, del_av;

  if (gauge != NULL)
    forallsites(i) su3_one( &gauge[i] );

  /* Do at most max_gauge_iter iterations, but stop after the second step if */
  /* the change in the avg gauge fixing action is smaller than gauge_fix_tol */

  old_av = del_av = 0;
  for (gauge_iter=0; gauge_iter < max_gauge_iter; gauge_iter++) {
    current_av = gaugefixstep(gauge_dir, relax_boost, gauge);
      
    if(gauge_iter != 0) {
      del_av = current_av - old_av;
      if (fabs(del_av) < gauge_fix_tol) break;
    }
    old_av = current_av;

    /* Reunitarize when iteration count is a multiple of REUNIT_INTERVAL */
    if((gauge_iter % REUNIT_INTERVAL) == (REUNIT_INTERVAL - 1)) {
/**	  node0_printf("step %d av gf action %.8e, delta %.3e\n",
		       gauge_iter,current_av,del_av); **/
      reunitarize(U);
    }
  }
  /* Reunitarize at the end, unless we just did it in the loop */
  if((gauge_iter % REUNIT_INTERVAL) != 0)
    reunitarize(U);

  if(this_node==0)
    printf("GFIX: Ended at step %d. Av gf action %.8e, delta %.3e\n",
	   gauge_iter,(double)current_av,(double)del_av);
}


/***********************************************************
 * and routines for doing gauge fixing on fields
 */

void vec_fix_gauge( su3_matrix *gauge, su3_vector *v, int parity )
{
  int i;
  su3_vector tv;

  /* vector <- u * vector */
  forparity(i,parity) {
    m_vec_mult( &gauge[i], &v[i], &tv );
    v[i] = tv;
  }
}

void ahmat_fix_gauge( su3_matrix *gauge, anti_hermitmat *a, int parity )
{
  int i;
  su3_matrix tm,tn;

  forparity(i,parity) {
    ahmat_uncompress( &a[i], &tm );
    m_mult_nn( &gauge[i], &tm, &tn );
    m_mult_na( &tn, &gauge[i], &tm );
    m_project_ahmat( &tm, &a[i] );
  }
}
