#include "./include/includes.h"
#include <float.h>

/* memory */
wilson_vector *cg_p, *cg_mp, *cg_res, *cg_temp;
float *cg_p_32, *cg_mp_32, *cg_res_32, *cg_temp_32;
int cg_ishift;

void malloc_cg64()
{
    MEMALIGN(cg_p,wilson_vector,sites_on_node);
    MEMALIGN(cg_mp,wilson_vector,sites_on_node);
    MEMALIGN(cg_res,wilson_vector,sites_on_node);
    MEMALIGN(cg_temp,wilson_vector,sites_on_node);
}
void free_cg64()
{
    FREE(cg_p,wilson_vector,sites_on_node);
    FREE(cg_mp,wilson_vector,sites_on_node);
    FREE(cg_res,wilson_vector,sites_on_node);
    FREE(cg_temp,wilson_vector,sites_on_node);
}
void malloc_cg32()
{
    MEMALIGN(cg_p_32,float,24*sites_on_node);
    MEMALIGN(cg_mp_32,float,24*sites_on_node);
    MEMALIGN(cg_res_32,float,24*sites_on_node);
    MEMALIGN(cg_temp_32,float,24*sites_on_node);
}
void free_cg32()
{
    FREE(cg_p_32,float,24*sites_on_node);
    FREE(cg_mp_32,float,24*sites_on_node);
    FREE(cg_res_32,float,24*sites_on_node);
    FREE(cg_temp_32,float,24*sites_on_node);
}

/* 32 bit CG. starts from 0 */
/* wvec: start contains src, end contains dest */
/* cg_p_32: search direction has to be set! */
double global_res;
int congrad_32(float *wvec, double myrsqmin, int maxniter, double shift)
{ 
    int it = 0, i, j;		/*counter for iterations */
    site *s;
    double flp=0;
    double alpha, mybeta;

    double source_norm;
    double rsq, oldrsq, pkp;

    double mytime;
    mytime = -dclock(  );

    malloc_cg32();
    
    /* r=res_32, rsq=|r|^2, source_norm=rsq */
    rsq = 0;
    FORALLSITES( i, s )
    {
        for (j=0;j<24;j++)
        {
            cg_p_32[24*i+j]=cg_res_32[24*i+j]=wvec[24*i+j];
            rsq+=(double)(cg_res_32[24*i+j]*cg_res_32[24*i+j]);
            wvec[24*i+j]=0.0;
        }
    }
    g_doublesum_KE( &rsq );
    source_norm=rsq;
    if( rsq <= myrsqmin )
    {
	goto end;
    }
    for ( it = 0; it < maxniter; it++ )
    {
      /* rsq -> oldrsq */
	oldrsq = rsq;

	/*  mp = M+M*p, pkp = p*mp  */
	multiply_fmat_32( cg_p_32, cg_temp_32, 1 );
	multiply_fmat_32( cg_temp_32, cg_mp_32, -1 );

        if (cg_ishift)
        {
            latutil_axpy_32((float)shift,cg_p_32,cg_mp_32,EVENANDODD);
        }
        pkp=latutil_rdot_32(cg_p_32,cg_mp_32,EVENANDODD);
	g_doublesum_KE( &pkp );

	/*  mybeta = rsq/pkp  */
	mybeta = rsq / pkp;

	/*  dest += mybeta*p */ 
        latutil_axpy_32((float)mybeta,cg_p_32,wvec,EVENANDODD);

	/*  r -=  mybeta*mp  */
        rsq=latutil_axpy_nrm2_32(-(float)mybeta,cg_mp_32,cg_res_32,EVENANDODD);
	g_doublesum_KE( &rsq );

	/* alpha = rsq/oldrsq    */
	alpha = rsq / oldrsq;

	/*  p = r+alpha*p */
        latutil_xpay_32(cg_res_32,(float)alpha,cg_p_32,EVENANDODD);

/* do not end the iteration */
/*         if( rsq <= source_norm*myrsqmin ) */
/* 	    goto end; */

	if (it%200==0) 
            node0_fprintf(file_o1,"congrad_32: %d prec= %e\n",
			  it,sqrt(rsq/source_norm));

    }
    verbose_fprintf( file_o1, "congrad_32: maxiter= %d reached prec= %e\n",
            maxniter,sqrt(rsq/source_norm));

end:
    mytime += dclock(  );
    /* flops in the infinit vol limit */
    flp = 2.0*1320.0+ 
	48.0+     
	24.0+     
	48.0+     
	72.0+     
	48.0;     
    if (cg_ishift) 
	flp += 48.0;
    flp *= 1.0*it;
    flp += 24.0;
    flp *= 1.0*sites_on_node;
    verbose_fprintf( file_o1, "congrad_32: it= %d\t%.3g sec\t%.3g GFlop/s/thread (%e Flop)\n",
            it, mytime, flp / mytime / 1.e9, flp );

    free_cg32();
    return it;
} 


/* original cg */
int congrad_64( wilson_vector *src, wilson_vector *dest, int maxniter, double myrsqmin, double shift)
{
    int it = 0, i;		/*counter for iterations */
    site *s;
    double alpha, mybeta;

    double source_norm;
    double rsqstop;
    double rsq, oldrsq, pkp;
    double flp;

    double mytime;
    mytime = -dclock(  );
    int nmatmul=0;

    malloc_cg64();
    
   
    /* shifted version */
    double absshift;
    cg_ishift=0;
    absshift=(shift>=0) ? (shift) : (-shift);
    if (absshift>DBL_EPSILON)
    {
        cg_ishift=1;
    }
       
    multiply_fmat( dest, cg_temp, 1 );
    multiply_fmat( cg_temp, cg_mp, -1 );

    /*r=p=src-(M+M)*dest, rsq=|r|^2, source_norm=|src|^2 */
    source_norm = rsq = 0;
    FORALLSITES( i, s )
    {
        if (cg_ishift)
            scalar_mult_add_wvec( &( cg_mp[i] ), &( dest[i] ), shift, &( cg_mp[i] ) );

        sub_wilson_vector( &( src[i] ), &( cg_mp[i] ), &( cg_res[i] ) );
        copy_wvec( &( cg_res[i] ), &( cg_p[i] ) );

        rsq += ( double ) magsq_wvec( &( cg_res[i] ) );
        source_norm += ( double ) magsq_wvec( &( src[i] ) );
    }
    g_doublesum_KE( &rsq );
    g_doublesum_KE( &source_norm );

/* do not end the iteration */
    rsqstop = myrsqmin * source_norm; 
/*     if( rsq <= rsqstop || source_norm <= myrsqmin ) */
/*     { */
/*         goto end; */
/*     } */

    it=0;
    while (/*rsq>rsqstop &&*/ (nmatmul)<2*maxniter)
    {
        /* // rsq -> oldrsq  */
        oldrsq = rsq;

        /* // mp = M+M*p, pkp = p*mp  */
        multiply_fmat( cg_p, cg_temp, 1 );
        multiply_fmat( cg_temp, cg_mp, -1 );
	nmatmul+=2;

        pkp = 0.0;
        if (cg_ishift)
        {
            FORALLSITES( i, s )
            {
                scalar_mult_add_wvec( &( cg_mp[i] ), &( cg_p[i] ), shift, &( cg_mp[i] ) );
                pkp += ( double ) wvec_rdot( &( cg_p[i] ), &( cg_mp[i] ) );
            }
        }
        else
        {
            FORALLSITES( i, s )
            {
                pkp += ( double ) wvec_rdot( &( cg_p[i] ), &( cg_mp[i] ) );
            }
        }
        g_doublesum_KE( &pkp );
	mybeta = rsq / pkp;

        /* // dest += mybeta*p  */
        FORALLSITES( i, s )
        {
            scalar_mult_add_wvec( &( dest[i] ), &( cg_p[i] ), mybeta, &( dest[i] ) );
        }

        /* // r -=  mybeta*mp  */
        rsq = 0.;
        FORALLSITES( i, s )
        {
            scalar_mult_add_wvec( &( cg_res[i] ), &( cg_mp[i] ), -mybeta, &( cg_res[i] ) );
            rsq += ( double ) magsq_wvec( &( cg_res[i] ) );
        }
        g_doublesum_KE( &rsq );

        alpha = rsq / oldrsq;

        /* // p = r+alpha*p */
        FORALLSITES( i, s )
        {
            scalar_mult_add_wvec( &( cg_res[i] ), &( cg_p[i] ), alpha, &( cg_p[i] ) );
        }
        if (it%200==0) 
            node0_fprintf(file_o1,"congrad_orig: %d prec= %e\n",
                    it,sqrt(rsq/source_norm));
                

        it++;
    }

    if ((nmatmul)>=2*maxniter)
        node0_fprintf( file_o1, "WARNING congrad_orig: not converged after it=%i: mvm= %d %e > %e\n", 
                it, nmatmul, sqrt( rsq / source_norm ), sqrt( myrsqmin ) );

end:
    if( source_norm <= myrsqmin )
    {
        FORALLSITES( i, s ) clear_wvec( &dest[i] );
    }

    /* timing */
    mytime += dclock(  );

    if (!cg_ishift)
        node0_fprintf(file_o1,"congrad_orig: end %d prec= %e mvm= %d time= %.3g\n",
           it,sqrt(rsq/source_norm),nmatmul,mytime);

    /* // flops in the infinit vol limit */
    flp = 2.0*1320.0  
	+24.0     
	+48.0     
	+48.0    
	+36.0    
	+48.0;   
    if (cg_ishift) 
	flp += 48.0; 
    flp *= 1.0*it;  
    flp += 2.0*1320.0
	+ 48.0    
	+ 24.0    
	+ 36.0    
	+ 36.0;   
    flp *= 1.0*sites_on_node;
    verbose_fprintf( file_o1, "congrad_orig: it= %d\t%.3g sec\t%.3g GFlop/s/thread (%e Flop)\n",
            it, mytime, flp / mytime / 1.e9, flp );

    free_cg64();

    return nmatmul;
}

