/*******************************************************************************
*
* File cg_iter.c
*
* Copyright (C) 2008 Bjorn Leder 2016 Jacob Finkenrath
*
* This software is distributed under the terms of the GNU General Public
* License (GPL)
*
* Based on QCDpbm-1.1 (http://luscher.web.cern.ch/luscher/QCDpbm/index.html)
*
*******************************************************************************/

#define CG_ITER_C

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "mpi.h"
#include "su3.h"
#include "su3fcts.h"
#include "random.h"
#include "lattice.h"
#include "linalg.h"
#include "sw_term.h"
#include "dirac.h"
#include "sflds.h"
#include "flags.h"
#include "uflds.h"
#include "utils.h"
#include "global.h"

spinor **ps;

static double wt_norm_square(int nflds)
{
   int my_rank,nmax,n,i,ib;
   /*float r;*/
   double wt1,wt2,wdt,wtav;

   MPI_Comm_rank(MPI_COMM_WORLD,&my_rank);   
   nmax=1;

   for (ib=0;ib<1;nmax*=2)
   {
      MPI_Barrier(MPI_COMM_WORLD);
      wt1=MPI_Wtime();

      for (n=0;n<nmax;n++)
      {
         for (i=0;i<nflds;i++)
            norm_square(VOLUME,1,ps[i]);
      }
      
      wt2=MPI_Wtime();
      wdt=wt2-wt1;

      MPI_Reduce(&wdt,&wtav,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);      

      if (my_rank==0)
      {
         wtav/=(double)(NPROC);

         if (wtav>2.0)
            ib=1;

         wtav/=(double)(nmax*nflds);
      }

      MPI_Bcast(&ib,1,MPI_INT,0,MPI_COMM_WORLD);
   }

   MPI_Bcast(&wtav,1,MPI_DOUBLE,0,MPI_COMM_WORLD);      
   
   return wtav;
}


static double wt_mulc_spinor_add(int nflds)
{
   int my_rank,nmax,n,i,ib;
   complex z;
   double wt1,wt2,wdt,wtav;

   MPI_Comm_rank(MPI_COMM_WORLD,&my_rank);   
   z.re=0.123f;
   z.im=0.456f;
   nmax=1;
   
   for (ib=0;ib<1;nmax*=2)
   {   
      MPI_Barrier(MPI_COMM_WORLD);
      wt1=MPI_Wtime();

      for (n=0;n<nmax;n++)
      {
         for (i=0;i<nflds;i+=2)
            mulc_spinor_add(VOLUME,ps[i],ps[i+1],z);

         z.re-=z.re;
         z.im-=z.im;
      }
      
      wt2=MPI_Wtime();
      wdt=wt2-wt1;

      MPI_Reduce(&wdt,&wtav,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);      

      if (my_rank==0)
      {
         wtav/=(double)(NPROC);

         if (wtav>2.0)
            ib=1;

         wtav/=(double)((nmax*nflds)/2);
      }

      MPI_Bcast(&ib,1,MPI_INT,0,MPI_COMM_WORLD);
   }

   MPI_Bcast(&wtav,1,MPI_DOUBLE,0,MPI_COMM_WORLD);
   
   return wtav;
}


static double wt_Qhat(int nflds)
{
   int my_rank,nmax,n,i,ib;
   double wt1,wt2,wdt,wtav;

   MPI_Comm_rank(MPI_COMM_WORLD,&my_rank);   
   nmax=1;
   
   for (ib=0;ib<1;nmax*=2)
   { 
      MPI_Barrier(MPI_COMM_WORLD);
      wt1=MPI_Wtime();

      for (n=0;n<nmax;n++)
      {
         for (i=0;i<nflds;i+=2)
            Dwhat(0.0,ps[i],ps[i+1]);
      }
      
      wt2=MPI_Wtime();
      wdt=wt2-wt1;

      MPI_Reduce(&wdt,&wtav,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);      

      if (my_rank==0)
      {
         wtav/=(double)(NPROC);

         if (wtav>2.0)
            ib=1;

         wtav/=(double)((nmax*nflds)/2);
      }
      

      MPI_Bcast(&ib,1,MPI_INT,0,MPI_COMM_WORLD);
   }

   MPI_Bcast(&wtav,1,MPI_DOUBLE,0,MPI_COMM_WORLD);
   
   return wtav;
}


void time_cg_iter(FILE *flog, double *wdt)
{
   int my_rank,nflds,n;
   double phi[2];
   double wdt0,wdt1,wdt2;
   
   MPI_Comm_rank(MPI_COMM_WORLD,&my_rank);

   if (my_rank==0)
   {
      printf("\n");
      printf("Single-precision data and programs\n");
      printf("-------------------------------\n\n");
      
      if ((VOLUME*sizeof(float))<(64*1024))
      {      
         printf("The local size of the gauge field is %d KB\n",
                (int)((72*VOLUME*sizeof(float))/(1024)));
         printf("The local size of a quark field is %d KB\n",
                (int)((24*VOLUME*sizeof(float))/(1024)));
      }
      else
      {
         printf("The local size of the gauge field is %d MB\n",
                (int)((72*VOLUME*sizeof(float))/(1024*1024)));
         printf("The local size of a quark field is %d MB\n",
                (int)((24*VOLUME*sizeof(float))/(1024*1024)));
      }
      
      printf("\n");
   }

   start_ranlux(0,12);
   phi[0]=0.0;
   phi[1]=0.0;
   set_bc_parms(3,1.0,1.0,1.0,1.0,phi,phi);
   
   set_lat_parms(5.5,1.0,0,NULL,1.978);
   print_lat_parms();
   set_sw_parms(-0.0123);
   geometry();
   /*alloc_u();
   alloc_ud();
   alloc_sw();
   alloc_swd();*/

   random_ud();
   chs_ubnd(-1);
   assign_ud2u();  
   
   sw_term(ODD_PTS);
   /*error(invert_swd(ODD_PTS)!=0,1,"main [time1.c]",
         "Inversion of swd on the odd sites was not safe");*/
   assign_swd2sw();

   nflds=(int)((4*1024*1024)/(VOLUME*sizeof(float)))+1;
   if ((nflds%2)==1)
      nflds+=1;
   alloc_ws(nflds);
   
   ps=reserve_ws(nflds);
   for (n=0;n<nflds;n++)
      random_s(VOLUME,ps[n],1.0f);

   error_chk();   
   wdt0=1.0e6*wt_norm_square(nflds)/(double)(VOLUME);

   if (my_rank==0)
   {
      printf("Program norm_square:\n");
      printf("Time per lattice point: %4.3f micro sec",wdt0);
      printf(" (%d Mflops/process)\n\n",(int)(48.0/wdt0));
      fflush(flog);
   }
   
   wdt1=1.0e6*wt_mulc_spinor_add(nflds)/(double)(VOLUME);   

   if (my_rank==0)
   {
      printf("Program mulc_spinor_add:\n");
      printf("Time per lattice point: %4.3f micro sec",wdt1);
      printf(" (%d Mflops/process)\n\n",(int)(96.0/wdt1));
      fflush(flog);
   }

   wdt2=1.0e6*wt_Qhat(nflds)/(double)(VOLUME);   

   if (my_rank==0)
   {
      printf("Program Dhat:\n");
      printf("Time per lattice point: %4.3f micro sec",wdt2);
      printf(" (%d Mflops/process)\n\n",(int)(1908.0/wdt2));
      fflush(flog);
   }

   /*free_u();
   free_ud();
   free_sw();
   free_swd();*/
   release_ws();
   
   wdt[0]=wdt0;
   wdt[1]=wdt1;
   wdt[2]=wdt2;
}
