#include "../include/lass.h"

/**
 *
 * 	@file kdcg.c
 *
 * 	@brief LASs-DSSs kdcg routine.
 *
 * 	LASs-DSSs is a software package provided by:
 * 	Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
 *
 * 	@author Sandra Catalán sandra.catalan@bsc.es 
 * 	@author Pedro Valero-Lara pedro.valero@bsc.es 
 * 	@date 2019-05-30
 * 	@reviewer
 * 	@modified 
 *
 **/

/**
 *  
 *	@ingroup ISSs
 *   
 *	Solves a system of linera equations with the CSR format:
 *
 *		AX = B
 *
 *	where X and B are vectors, and 
 *  A is a sparse matrix stored in the next three vectors 
 *  VAL_A, COL_IND_A and ROW_PTR_A according to CSR format.
 *
 **/

/**
 *    
 *	@param[in]
 *  N   		int.
 *	    		N specifies the number of columns of A.
 *
 *	@param[in]
 *  VAL_A		double *.
 *	       		VAL_A is a pointer to a vector which specifies the non-zero 
 *          	values of the original matrix A. 
 *   
 *	@param[in]
 *  COL_IND_A  	unsigned int *.
 * 		    	COL_IND_A is a pointer to a vector which specifies the 
 *				column index of each value of A in the original matrix. 
 *   
 *	@param[in]
 *  ROW_PTR_A   unsigned int *.
 * 		    	ROW_PTR_A is a pointer to a vector which specifies the number of
 *				values that A contains in each row of the original matrix.
 *
 *	@param[in]
 *  B  			double *.
 *          	B is a pointer to a vector of dimension N. 
 *   
 *  @param[in, out]
 *  X       	double *.
 *          	X is a pointer to a vector of dimension N. 
 *				On exxit, X is overwritten by the CG result
 *
 *
 **/

/**
 *
 *	@retval Success sucessful exit
 *
 **/

/**
 * 
 * @sa isss_dcg
 *
 **/

enum LASS_RETURN kdcg ( int N, 
		const double *VAL_A, const int *ROW_PTR_A, const int *COL_IND_A, 
		const double *B,
		double *X )
{
	
    int i, j, end;
	double tolerance = 1.0e-10, one = 1.0, mOne = -1.0, zero = 0.0, alpha, beta,
		   bnorm, rnorm, alpha1, alpha2, beta1, beta2, condition;
	int chunks_per_core = 1;
	int number_of_cores = NUM_CORES;
    int granularity = N / (chunks_per_core * number_of_cores);
	double *r, *p, *rOld, *s, *rAux, *bAux;
	struct timeval endT, startT;

	r    = (double*) malloc( N * sizeof(double) );
	p    = (double*) malloc( N * sizeof(double) );
	s    = (double*) malloc( N * sizeof(double) );
	rOld = (double*) malloc( N * sizeof(double) );
	rAux = (double*) malloc( N * sizeof(double) );
	bAux = (double*) malloc( N * sizeof(double) );

	memcpy( r, B, N * sizeof( double ) ); 
	memcpy( p, r, N * sizeof( double ) );

	i = 0;

	while ( i < N )
	{	
		#pragma omp taskloop grainsize(1) label(memcpy)
		for( j = 0; j < N; j += granularity )
		{
			end = j + MIN( granularity, N-j );
			memcpy( &rOld[j], &r[j], (end-j) * sizeof( double ) );
		}

		kdspmv( N, N, one, VAL_A, ROW_PTR_A, COL_IND_A, p, zero, s );

		alpha1 = 0.0;
		#pragma omp taskloop grainsize(1) reduction(+: alpha1) label(ddot1)
		for( j = 0; j < N; j += granularity )
		{
			end = j + MIN( granularity, N-j );
			alpha1 += cblas_ddot( end-j, &r[j], 1, &r[j], 1 );
		}	

		alpha2=0.0;
		#pragma omp taskloop grainsize(1) reduction(+: alpha2) label(ddot2)
        for( j = 0; j < N; j += granularity )
        {
            end = j + MIN( granularity, N-j );
            alpha2 += cblas_ddot( end - j, &p[j], 1, &s[j], 1 );
        }

		alpha = alpha1/alpha2;

		#pragma omp taskloop grainsize(1) label(axpy1)
		for( j = 0; j < N; j += granularity )
		{
			end = j + MIN(granularity, N-j);
			cblas_daxpy( end - j,  alpha, &p[j], 1, &X[j], 1);
		}
		
		#pragma omp taskloop grainsize(1) label(axpy2)
        for( j = 0; j < N; j += granularity )
        {
            end = j + MIN( granularity, N-j );
            cblas_daxpy( end - j, -alpha, &s[j], 1, &r[j], 1);
        }

        beta1 = 0.0;
		#pragma omp taskloop grainsize(1) reduction(+: beta1) label(ddot3)
        for( j = 0; j < N; j += granularity )
        {
            end = j + MIN( granularity, N-j );
            beta1 += cblas_ddot( end - j, &r[j], 1, &r[j], 1);
        }

        beta2 = 0.0;
		#pragma omp taskloop grainsize(1) reduction(+: beta2) label(ddot4)
        for( j = 0; j < N; j += granularity)
        {
            end = j + MIN( granularity, N-j );
            beta2 += cblas_ddot( end - j, &rOld[j], 1, &rOld[j], 1 );
        }

        beta = beta1/beta2;

		#pragma omp taskloop grainsize(1) label(axpy3)
        for( j = 0; j < N; j += granularity )
        {
            end = j + MIN( granularity, N-j );
			memcpy( &rAux[j], &r[j], (end - j) * sizeof( double ) );	
            cblas_daxpy( end-j, beta, &p[j], 1, &rAux[j], 1);
			memcpy( &p[j], &rAux[j], (end - j) * sizeof( double ) );
        }

        condition = 0.0;
		#pragma omp taskloop grainsize(1) reduction(+: condition) label(ddot5)
        for(j = 0; j < N; j += granularity)
        {
            end = j + MIN( granularity, N-j );
            condition += cblas_ddot( end-j, &r[j], 1, &r[j], 1 );
        }

		if ( sqrt( condition )  < tolerance)
			break;

		i++;
	}   	

	return Success;
}
