/*************************************************************************************

    Grid physics library, www.github.com/paboyle/Grid 

    Source file: ./lib/qcd/utils/WilsonLoops.h

    Copyright (C) 2015

Author: Azusa Yamaguchi <ayamaguc@staffmail.ed.ac.uk>
Author: Peter Boyle <paboyle@ph.ed.ac.uk>
Author: neo <cossu@post.kek.jp>
Author: paboyle <paboyle@ph.ed.ac.uk>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License along
    with this program; if not, write to the Free Software Foundation, Inc.,
    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

    See the full license in the file "LICENSE" in the top level distribution directory
*************************************************************************************/
/*  END LEGAL */
#ifndef QCD_UTILS_WILSON_LOOPS_H
#define QCD_UTILS_WILSON_LOOPS_H
namespace Grid {
  namespace QCD {
    
    // Common wilson loop observables
    template<class Gimpl>
      class WilsonLoops : public Gimpl {
    public:
      
      INHERIT_GIMPL_TYPES(Gimpl);
      
      typedef typename Gimpl::GaugeLinkField GaugeMat;
      typedef typename Gimpl::GaugeField     GaugeLorentz;
      
      //////////////////////////////////////////////////
      // directed plaquette oriented in mu,nu plane
      //////////////////////////////////////////////////
      static void dirPlaquette(GaugeMat &plaq,const std::vector<GaugeMat> &U, const int mu, const int nu)
      {
	// Annoyingly, must use either scope resolution to find dependent base class, 
	// or this-> ; there is no "this" in a static method. This forces explicit Gimpl scope
	// resolution throughout the usage in this file, and rather defeats the purpose of deriving
	// from Gimpl.
	plaq= Gimpl::CovShiftBackward(U[mu],mu,
				      Gimpl::CovShiftBackward(U[nu],nu,
							      Gimpl::CovShiftForward (U[mu],mu,U[nu])));
      }
      //////////////////////////////////////////////////
      // trace of directed plaquette oriented in mu,nu plane
      //////////////////////////////////////////////////
      static void traceDirPlaquette(LatticeComplex &plaq, const std::vector<GaugeMat> &U, const int mu, const int nu)
      {
	GaugeMat sp(U[0]._grid);
	dirPlaquette(sp,U,mu,nu);
	plaq=trace(sp);
      }
      //////////////////////////////////////////////////
      // sum over all planes of plaquette
      //////////////////////////////////////////////////
      static void sitePlaquette(LatticeComplex &Plaq,const std::vector<GaugeMat> &U)
      {
	LatticeComplex sitePlaq(U[0]._grid);
	Plaq=zero;
	for(int mu=1;mu<Nd;mu++){
	  for(int nu=0;nu<mu;nu++){
	    traceDirPlaquette(sitePlaq,U,mu,nu);
	    Plaq = Plaq + sitePlaq;
	  }
	}
      }
      //////////////////////////////////////////////////
      // sum over all x,y,z,t and over all planes of plaquette
      //////////////////////////////////////////////////
      static RealD sumPlaquette(const GaugeLorentz &Umu){
	std::vector<GaugeMat> U(4,Umu._grid);
	
	for(int mu=0;mu<Nd;mu++){
	  U[mu] = PeekIndex<LorentzIndex>(Umu,mu);
	}
	
	LatticeComplex Plaq(Umu._grid);
	
	sitePlaquette(Plaq,U);
	
	TComplex Tp = sum(Plaq);
	Complex p  = TensorRemove(Tp);
	return p.real();
      }
      //////////////////////////////////////////////////
      // average over all x,y,z,t and over all planes of plaquette
      //////////////////////////////////////////////////
      static RealD avgPlaquette(const GaugeLorentz &Umu){
	
	RealD sumplaq = sumPlaquette(Umu);
	
	double vol = Umu._grid->gSites();
	
	double faces = (1.0*Nd*(Nd-1))/2.0;
	
	return sumplaq/vol/faces/Nc; // Nd , Nc dependent... FIXME
      }
      static RealD linkTrace(const GaugeLorentz &Umu){
	std::vector<GaugeMat> U(4,Umu._grid);
	
	LatticeComplex Tr(Umu._grid); Tr=zero;
	for(int mu=0;mu<Nd;mu++){
	  U[mu] = PeekIndex<LorentzIndex>(Umu,mu);
	  Tr = Tr+trace(U[mu]);
	}
	
	TComplex Tp = sum(Tr);
	Complex p  = TensorRemove(Tp);
	
	double vol = Umu._grid->gSites();
	
	return p.real()/vol/4.0/3.0;
      };
      //////////////////////////////////////////////////
      // the sum over all staples on each site
      //////////////////////////////////////////////////
      static void Staple(GaugeMat &staple,const GaugeLorentz &Umu,int mu){
	
	GridBase *grid = Umu._grid;
	
	std::vector<GaugeMat> U(4,grid);
	for(int d=0;d<Nd;d++){
	  U[d] = PeekIndex<LorentzIndex>(Umu,d);
	}
	staple = zero;
		
	
	for(int nu=0;nu<Nd;nu++){
	  
	  if(nu != mu) {
	    
	    // mu
	    // ^
	    // |__>  nu
	    
	    //    __ 
	    //      |
	    //    __|
	    //
	    
	    staple+=Gimpl::ShiftStaple(Gimpl::CovShiftForward (U[nu],nu, 
							       Gimpl::CovShiftBackward(U[mu],mu,
										       Gimpl::CovShiftIdentityBackward(U[nu],nu))),mu);
	    
	    //  __ 
	    // |   
	    // |__ 
	    //
	    //
	    staple+=Gimpl::ShiftStaple(Gimpl::CovShiftBackward(U[nu],nu,		  		  
							       Gimpl::CovShiftBackward(U[mu],mu,U[nu])),mu);
	  }
	}
      }
      
      //////////////////////////////////////////////////
      // the sum over all staples on each site in direction mu,nu
      //////////////////////////////////////////////////
      static void Staple(GaugeMat &staple,const GaugeLorentz &Umu,int mu, int nu){

	GridBase *grid = Umu._grid;
	
	std::vector<GaugeMat> U(4,grid);
	for(int d=0;d<Nd;d++){
	  U[d] = PeekIndex<LorentzIndex>(Umu,d);
	}
	staple = zero;
	
	
	if(nu != mu) {
	  
	  // mu
	  // ^
	  // |__>  nu
	  
	  //    __ 
	  //      |
	  //    __|
	  //

	  staple+=Gimpl::ShiftStaple(
				     Gimpl::CovShiftForward (U[nu],nu, 
							     Gimpl::CovShiftBackward(U[mu],mu,
										     Gimpl::CovShiftIdentityBackward(U[nu],nu))),mu);

	  //  __ 
	  // |   
	  // |__ 
	  //
	  //
	  staple+=Gimpl::ShiftStaple(  
				     Gimpl::CovShiftBackward(U[nu],nu,		  		  
							     Gimpl::CovShiftBackward(U[mu],mu,U[nu])),mu);
      
	}
      }


      //////////////////////////////////////////////////
      // the sum over all staples on each site in direction mu,nu, upper part
      //////////////////////////////////////////////////
      static void StapleUpper(GaugeMat &staple,const GaugeLorentz &Umu,int mu, int nu){

	GridBase *grid = Umu._grid;

	std::vector<GaugeMat> U(4,grid);
	for(int d=0;d<Nd;d++){
	  U[d] = PeekIndex<LorentzIndex>(Umu,d);
	}
	staple = zero;


    
	if(nu != mu) {
      
	  // mu
	  // ^
	  // |__>  nu
      
	  //    __ 
	  //      |
	  //    __|
	  //

	  staple+=Gimpl::ShiftStaple(
				     Gimpl::CovShiftForward (U[nu],nu, 
							     Gimpl::CovShiftBackward(U[mu],mu,
										     Gimpl::CovShiftIdentityBackward(U[nu],nu))),mu);

	}
      }






      //////////////////////////////////////////////////////
      // Similar to above for rectangle is required
      //////////////////////////////////////////////////////
      static void dirRectangle(GaugeMat &rect,const std::vector<GaugeMat> &U, const int mu, const int nu)
      {
	rect =  Gimpl::CovShiftForward(U[mu],mu,Gimpl::CovShiftForward(U[mu],mu,U[nu]))* // ->->|
	  adj(Gimpl::CovShiftForward(U[nu],nu,Gimpl::CovShiftForward(U[mu],mu,U[mu]))) ;
	rect = rect + 
          Gimpl::CovShiftForward(U[mu],mu,Gimpl::CovShiftForward(U[nu],nu,U[nu]))* // ->||
	  adj(Gimpl::CovShiftForward(U[nu],nu,Gimpl::CovShiftForward(U[nu],nu,U[mu]))) ;
      }
      static void traceDirRectangle(LatticeComplex &rect, const std::vector<GaugeMat> &U, const int mu, const int nu)
      {
	GaugeMat sp(U[0]._grid);
	dirRectangle(sp,U,mu,nu);
	rect=trace(sp);
      }
      static void siteRectangle(LatticeComplex &Rect,const std::vector<GaugeMat> &U)
      {
	LatticeComplex siteRect(U[0]._grid);
	Rect=zero;
	for(int mu=1;mu<Nd;mu++){
	  for(int nu=0;nu<mu;nu++){
	    traceDirRectangle(siteRect,U,mu,nu);
	    Rect = Rect + siteRect;
	  }
	}
      }
      //////////////////////////////////////////////////
      // sum over all x,y,z,t and over all planes of plaquette
      //////////////////////////////////////////////////
      static RealD sumRectangle(const GaugeLorentz &Umu){
	std::vector<GaugeMat> U(4,Umu._grid);

	for(int mu=0;mu<Nd;mu++){
	  U[mu] = PeekIndex<LorentzIndex>(Umu,mu);
	}

	LatticeComplex Rect(Umu._grid);
    
	siteRectangle(Rect,U);
    
	TComplex Tp = sum(Rect);
	Complex p  = TensorRemove(Tp);
	return p.real();
      }
      //////////////////////////////////////////////////
      // average over all x,y,z,t and over all planes of plaquette
      //////////////////////////////////////////////////
      static RealD avgRectangle(const GaugeLorentz &Umu){

	RealD sumrect = sumRectangle(Umu);
    
	double vol = Umu._grid->gSites();
    
	double faces = (1.0*Nd*(Nd-1)); // 2 distinct orientations summed
    
	return sumrect/vol/faces/Nc; // Nd , Nc dependent... FIXME
      }

      //////////////////////////////////////////////////
      // the sum over all staples on each site
      //////////////////////////////////////////////////
      static void RectStapleDouble(GaugeMat &U2,const GaugeMat & U,int mu){
	U2 = U * Cshift(U,mu,1);
      }

      ////////////////////////////////////////////////////////////////////////////
      // Hop by two optimisation strategy does not work nicely with Gparity. (could do,
      // but need to track two deep where cross boundary and apply a conjugation).
      // Must differentiate this in Gimpl, and use Gimpl::isPeriodicGaugeField to do so .
      ////////////////////////////////////////////////////////////////////////////
      static void RectStapleOptimised(GaugeMat &Stap,std::vector<GaugeMat> &U2,std::vector<GaugeMat> &U,int mu){

	Stap = zero;

	GridBase *grid = U[0]._grid;

	GaugeMat Staple2x1 (grid);
	GaugeMat tmp (grid);

	for(int nu=0;nu<Nd;nu++){
	  if ( nu!=mu) {

	    // Up staple    ___ ___ 
	    //             |       |
	    tmp = Cshift(adj(U[nu]),nu,-1); 
	    tmp = adj(U2[mu])*tmp;
	    tmp = Cshift(tmp,mu,-2);

	    Staple2x1 = Gimpl::CovShiftForward (U[nu],nu,tmp);


	    // Down staple
	    //             |___ ___|
	    //
	    tmp = adj(U2[mu])*U[nu];
	    Staple2x1+= Gimpl::CovShiftBackward(U[nu],nu,Cshift(tmp,mu,-2));


	    //              ___ ___
	    //             |    ___|
	    //             |___ ___|
	    //

	    Stap+= Cshift(Gimpl::CovShiftForward (U[mu],mu,Staple2x1),mu,1);

	    //              ___ ___
	    //             |___    |
	    //             |___ ___|
	    //

	    //	tmp= Staple2x1* Cshift(U[mu],mu,-2);
	    //	Stap+= Cshift(tmp,mu,1) ;
	    Stap+= Cshift(Staple2x1,mu,1)*Cshift(U[mu],mu,-1); ;

	    //       --    
	    //      |  |              
	    //          
	    //      |  | 
	
	    tmp = Cshift(adj(U2[nu]),nu,-2);
	    tmp = Gimpl::CovShiftBackward(U[mu],mu,tmp);
	    tmp = U2[nu]*Cshift(tmp,nu,2);
	    Stap+= Cshift(tmp, mu, 1);

	    //      |  |              
	    //          
	    //      |  | 
	    //       -- 
	
	    tmp = Gimpl::CovShiftBackward(U[mu],mu,U2[nu]);
	    tmp = adj(U2[nu])*tmp;
	    tmp = Cshift(tmp,nu,-2);
	    Stap+=Cshift(tmp, mu, 1);
	  }}


      }

      static void RectStaple(GaugeMat &Stap,const GaugeLorentz & Umu,int mu)
      {
	RectStapleUnoptimised(Stap,Umu,mu);
      }
      static void RectStaple(const GaugeLorentz & Umu,GaugeMat &Stap,
			     std::vector<GaugeMat> &U2,
			     std::vector<GaugeMat> &U, int mu)
      {
	if ( Gimpl::isPeriodicGaugeField() ){ 
	  RectStapleOptimised(Stap,U2,U,mu);
	} else {
	  RectStapleUnoptimised(Stap,Umu,mu);
	}
      }

      static void RectStapleUnoptimised(GaugeMat &Stap,const GaugeLorentz &Umu,int mu){
	GridBase *grid = Umu._grid;

	std::vector<GaugeMat> U(4,grid);
	for(int d=0;d<Nd;d++){
	  U[d] = PeekIndex<LorentzIndex>(Umu,d);
	}

	Stap=zero;

	for(int nu=0;nu<Nd;nu++){
	  if ( nu!=mu) {
	    //           __ ___ 
	    //          |    __ |
	    //
	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftForward (U[mu],mu,
							      Gimpl::CovShiftForward (U[nu],nu,
										      Gimpl::CovShiftBackward(U[mu],mu,
													      Gimpl::CovShiftBackward(U[mu],mu,
																      Gimpl::CovShiftIdentityBackward(U[nu],nu))))) , mu);

	    //              __ 
	    //          |__ __ |

	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftForward (U[mu],mu,
							      Gimpl::CovShiftBackward(U[nu],nu,
										      Gimpl::CovShiftBackward(U[mu],mu,
													      Gimpl::CovShiftBackward(U[mu],mu, U[nu])))) , mu);

	    //           __ 
	    //          |__ __ |

	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftBackward(U[nu],nu,
							      Gimpl::CovShiftBackward(U[mu],mu,
										      Gimpl::CovShiftBackward(U[mu],mu,
													      Gimpl::CovShiftForward(U[nu],nu,U[mu])))) , mu);

	    //           __ ___ 
	    //          |__    |

	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftForward (U[nu],nu,
							      Gimpl::CovShiftBackward(U[mu],mu,
										      Gimpl::CovShiftBackward(U[mu],mu,
													      Gimpl::CovShiftBackward(U[nu],nu,U[mu])))) , mu);

	    //       --    
	    //      |  |              
	    //          
	    //      |  | 
     
	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftForward(U[nu],nu,
							     Gimpl::CovShiftForward(U[nu],nu,
										    Gimpl::CovShiftBackward(U[mu],mu,
													    Gimpl::CovShiftBackward(U[nu],nu,
																    Gimpl::CovShiftIdentityBackward(U[nu],nu))))) , mu);


	    //      |  |              
	    //          
	    //      |  | 
	    //       -- 
     
	    Stap+= Gimpl::ShiftStaple(
				      Gimpl::CovShiftBackward(U[nu],nu,
							      Gimpl::CovShiftBackward(U[nu],nu,
										      Gimpl::CovShiftBackward(U[mu],mu,
													      Gimpl::CovShiftForward (U[nu],nu,U[nu])))) , mu);
	  }}
      }


    };


    typedef WilsonLoops<PeriodicGimplR> ColourWilsonLoops;
    typedef WilsonLoops<PeriodicGimplR> U1WilsonLoops;
    typedef WilsonLoops<PeriodicGimplR> SU2WilsonLoops;
    typedef WilsonLoops<PeriodicGimplR> SU3WilsonLoops;

  }}

#endif