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

/**
 *
 * 	@file kdtpgesv.c
 *
 * 	@brief LASs-DDSs kdtpgesv routine.
 *
 * 	LASs-DDSs is a software package provided by:
 * 	Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
 *
 * 	@author Pedro Valero-Lara pedro.valero@bsc.es
 * 	@date 2018-09-20
 * 	@reviewer 
 * 	@modified 
 *
 **/

/**
 *  
 *	@ingroup DDSS
 *   
 *	Solves a system of linear equations A \times X = B, where A is a N-by-N
 *	general matrix and X and B are N-by-NRHS matrices. The matrix A is 
 *	factorized using the LU descomposition with tiled-pivoting.
 *	The matrix A is descomposed as:
 *
 *		A = P * L * U
 *
 *	where P is a permutation matrix, L is a lower triangular matrix with unit 
 *  diagonal elements and U is an upper triangular matrix.
 *
**/

/**
 *
 *	@param[in]
 *	N       int.
 *          N specifies the order of the square matrix A. N >= 0.
 *	
 *	NRHS    int.
 *          NRHS specifies the number of right-hand-sides (number 
 *			of columns of B). NRHS >= 0.
 *
 *	@param[in,out]
 *  A  		double *.
 *	    	A is a pointer to a regular matrix of dimension N-by-LDA.
 *          On exit, if return value is Success, the matrix A is overwriten by 
 *          the factors L and U. The unit diagonal elements of L are not stored.
 *   
 *	@param[in]
 *  LDA     int.
 * 	    	LDA specifies the number of columns of A ( row-major order ).
 * 	    	LDA must be at least max( 1, N ).
 * 	    	
 * 	@param[out]
 * 	IPIV    int *.
 * 			ipiv is a pointer to an array of dimesion at least 
 *			max( 1, min ( M, N ) ). ipiv( i ) = j, 1 <= i <= min( M, N ) implies
 *			that rows i and j have been interchanged.    	
 * 
 *	@param[in,out]
 *  B  		double *.
 *	    	B is a pointer to a matrix of dimension N by NRHS, which stores the 
 *			right-hand-sides of the systems of linear equations.  
 *			(row-major order). On exit, if return value is Success, the matrix B
 * 			is overwriten by the solution matrix X.
 *   
 *	@param[in]
 *  LDB     int.
 * 	    	LDB specifies the number of columns of B ( row-major order ).
 * 	    	LDB must be at least max( 1, NRHS ).
 *
 **/

/**
 *
 *	@retval Success successful exit
 *	@retval NoSuccess unsuccessful exit
 *
 **/

/**
 *
 *	@sa ddss_tile
 *	@sa ddss_flat2tiled
 *	@sa ddss_tiled2flat
 *
 **/ 

enum LASS_RETURN kdtpgesv( int N, int NRHS, 
	double *A, int LDA, 
	int *IPIV, 
	double *B, int LDB )
{

	// Local variables
	int i;
	int nt, nrhst;
	int ni, mi, nni, ki, nrhsi;
	int tile_size_n;
	int tile_size_m;
	int tile_size_nn;
	int tile_size_nrhs;
	int tile_size_k;
	int local_ipiv_size;
	int current_position;
	enum LASS_RETURN ret;
	int	*ipiv_local;

	// Number of tiles
	if ( N % TILE_SIZE == 0 )
	{
		nt = N / TILE_SIZE;
	}
	else
	{
		nt = ( N / TILE_SIZE ) + 1;
	}

	if ( NRHS % TILE_SIZE == 0 )
	{
		nrhst = NRHS / TILE_SIZE;
	}
	else
	{
		nrhst = ( NRHS / TILE_SIZE ) + 1;
	}

	/***************************
	--Tile matrices allocation--
	***************************/
		
	double ( *TILE_A )[nt][TILE_SIZE * TILE_SIZE] = malloc ( nt * nt * 
		TILE_SIZE * TILE_SIZE * sizeof( double ) );
	
	if ( TILE_A == NULL )
	{
		fprintf( stderr, "Failure in kdtpgesv for matrix TILE_A\n" );
		return NoSuccess;
	}
	
	double ( *TILE_B )[nrhst][TILE_SIZE * TILE_SIZE] = malloc(
			 nt * nrhst * TILE_SIZE * TILE_SIZE * sizeof( double ) );

	if ( TILE_B == NULL ) 
	{
		fprintf( stderr, "Failure in kdtpgesv for matrix TILE_B\n" );
		return NoSuccess;
	}

	ipiv_local = ( int* ) malloc( N * sizeof( int ) );
	
	if ( ipiv_local == NULL )
	{
		fprintf( stderr, "Failure in kdtpgesv for ipiv_local\n" );
		return NoSuccess;
	}

	/************************************************
	// --From flat data layout to tiled data layout--	
	************************************************/
	
	// From flat matrix A to tile matrix TILE_A
	ddss_dflat2tiled( N, N, A, LDA, nt, nt, TILE_A );

	// From flat matrix B to tile matrix TILE_B 
	ddss_dflat2tiled( N, NRHS, B, LDB, nt, nrhst, TILE_B );
	
	/****************
	--DGETRF tile--
	****************/	
	
	current_position = 0;
	i = 0;
	for ( ki = 0; ki < nt; ki++ )
	{
		tile_size_k = ddss_tile_size( N, ki );
		
		#pragma oss task inout( TILE_A[ki][ki] ) \
			inout( ipiv_local ) \
			inout( current_position ) \
			out( local_ipiv_size ) \
			shared( TILE_A ) \
			firstprivate( ki, i ) \
			label( dgetrf )
		{
			LAPACKE_dgetrf( CblasRowMajor,
							tile_size_k, tile_size_k, 
							TILE_A[ki][ki], tile_size_k,
							ipiv_local + current_position );
			
			// Update the global ipiv array
			local_ipiv_size = tile_size_k;
			for ( i = 0; i < local_ipiv_size; i++ )
			{
				( IPIV + current_position )[i] =
						( ipiv_local + current_position )[i] + current_position;
			}
			current_position += local_ipiv_size;
		}
		for ( mi = ki + 1; mi < nt; mi++ )
		{
			tile_size_m = ddss_tile_size( N, mi );

			#pragma oss task in( TILE_A[ki][ki] ) \
				inout(TILE_A[mi][ki]) \
				shared( TILE_A ) \
				firstprivate( mi, ki ) \
				label( dtrsm_below )
			cblas_dtrsm( CblasRowMajor, 
				     	 ( CBLAS_SIDE ) Right, ( CBLAS_UPLO ) Upper, 
				     	 ( CBLAS_TRANSPOSE ) NoTrans, ( CBLAS_DIAG ) NonUnit, 
					 	 tile_size_m, tile_size_k, 
					 	 1.0, TILE_A[ki][ki], tile_size_k, 
						  	  TILE_A[mi][ki], tile_size_k );
		}
		for ( ni = 0; ni < ki; ni++ )
		{		
			tile_size_n = ddss_tile_size( N, ni );
			// Swap the tiles to the left of A[ki][ki]
			#pragma oss task in ( ipiv_local ) \
				in( current_position ) \
				in( local_ipiv_size ) \
				inout( TILE_A[ki][ni] ) \
				shared( TILE_A ) \
				firstprivate( ki, ni ) \
				label( dlaswp_left )
			LAPACKE_dlaswp( CblasRowMajor,
							tile_size_n, TILE_A[ki][ni], tile_size_n,
							1, local_ipiv_size,
							ipiv_local + current_position - local_ipiv_size, 
							1 );
			
		}
		for ( nrhsi = 0; nrhsi < nrhst; nrhsi++ )
		{		
			tile_size_nrhs = ddss_tile_size( NRHS, nrhsi );
			// Swap on the B[ki][nrhsi] tile
			#pragma oss task in ( ipiv_local ) \
				in( current_position ) \
				in( local_ipiv_size ) \
				inout( TILE_B[ki][nrhsi] ) \
				shared( TILE_B ) \
				firstprivate( ki, nrhsi ) \
				label( dlaswp_on_B )
			LAPACKE_dlaswp( CblasRowMajor,
							tile_size_nrhs, TILE_B[ki][nrhsi], tile_size_nrhs,
							1, local_ipiv_size,
							ipiv_local + current_position - local_ipiv_size, 
							1 );
			
		}
		for ( ni = ki + 1; ni < nt; ni++ )
		{
			tile_size_n = ddss_tile_size( N, ni );
			// Swap the tiles to the right of A[ki][ki]
			#pragma oss task in ( ipiv_local ) \
				in( current_position ) \
				in( local_ipiv_size ) \
				inout( TILE_A[ki][ni] ) \
				shared( TILE_A ) \
				firstprivate( ki, ni ) \
				label( dlaswp_right )
			LAPACKE_dlaswp( CblasRowMajor,
							tile_size_n, TILE_A[ki][ni], tile_size_n,
							1, local_ipiv_size,
							ipiv_local + current_position - local_ipiv_size, 
							1 );

			#pragma oss task in( TILE_A[ki][ki] ) \
				inout(TILE_A[ki][ni]) \
				shared( TILE_A ) \
				firstprivate( ni, ki ) \
				label( dtrsm_right )
			cblas_dtrsm( CblasRowMajor, 
				     	 ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower, 
				     	 ( CBLAS_TRANSPOSE ) NoTrans, ( CBLAS_DIAG ) Unit, 
					 	 tile_size_k, tile_size_n, 
					 	 1.0, TILE_A[ki][ki], tile_size_k, 
						  	  TILE_A[ki][ni], tile_size_n );
				
			for ( nni = ki + 1; nni < nt; nni++ )
			{
				tile_size_nn = ddss_tile_size( N, nni );
		
				#pragma oss task in( TILE_A[nni][ki] ) \
					in( TILE_A[ki][ni] ) \
					inout( TILE_A[nni][ni] ) \
					shared( TILE_A ) \
					firstprivate( ni, nni, ki ) \
					label( dgemm )
				cblas_dgemm( CblasRowMajor, 
							 ( CBLAS_TRANSPOSE ) NoTrans,
							 ( CBLAS_TRANSPOSE ) NoTrans,
							 tile_size_nn, 
							 tile_size_n, 
							 TILE_SIZE,
							 -1.0, TILE_A[nni][ki], tile_size_k,
							 	   TILE_A[ki][ni],  tile_size_n,
							  1.0, TILE_A[nni][ni], tile_size_n );			
			}
		}
	}
		
	// --From tile data layout to flat data layout--	
	ddss_dtiled2flat( N, N, A, LDA, nt, nt, TILE_A );

	// Triandular solve
	// --SIDE = Left & UPLO = Lower & TRANS_A = NoTrans & Unit-- 
	for ( ki = 0; ki < nt; ki++ )
	{
		tile_size_k = ddss_tile_size( N, ki );

		for ( ni = 0; ni < nrhst; ni++ )
		{
			tile_size_n = ddss_tile_size( NRHS, ni );
						
			#pragma oss task in( TILE_A[ki][ki] ) \
				inout(  TILE_B[ki][ni] ) \
				shared( TILE_A, TILE_B ) \
				firstprivate( ki, ni ) \
				label( dtrsm_dtrsm1 )
			cblas_dtrsm( CblasRowMajor,
						 ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower,
						 ( CBLAS_TRANSPOSE ) NoTrans,
						 ( CBLAS_DIAG ) Unit,
						 tile_size_k,
						 tile_size_n,
						 1.0, TILE_A[ki][ki], tile_size_k,
							  TILE_B[ki][ni], tile_size_n );
		}
		for ( nni = 0; nni < nrhst; nni++ ) 
		{
			tile_size_nn = ddss_tile_size( NRHS, nni );
				
			for ( mi = ki + 1; mi < nt; mi++ )
			{
					
				#pragma oss task in( TILE_A[mi][ki] ) \
                    in( TILE_B[ki][nni] ) \
                    inout( TILE_B[mi][nni] ) \
                    shared( TILE_A, TILE_B ) \
                    firstprivate( ki, mi, nni ) \
                    label( dtrsm_dgemm1 )
				cblas_dgemm( CblasRowMajor,
							 CblasNoTrans, CblasNoTrans,
							 tile_size_k,
							 tile_size_nn,
							 tile_size_k,
							 -1.0,  TILE_A[mi][ki],  tile_size_k,
								   TILE_B[ki][nni], tile_size_nn,
							  1.0, TILE_B[mi][nni], tile_size_nn );
			}
		}
	}

	// Triandular solve
	// --SIDE = Left & UPLO = Upper & TRANS_A = NoTrans & NonUnit--
	for ( ki = nt - 1; ki >= 0; ki-- ) 
	{
		tile_size_k = ddss_tile_size( N, ki );
		
		for ( ni = 0; ni < nrhst; ni++ ) 
		{
			tile_size_n = ddss_tile_size( NRHS, ni );

			#pragma oss task in( TILE_A[ki][ki] ) \
				inout( TILE_B[ki][ni] ) \
				shared( TILE_A, TILE_B ) \
				firstprivate( ki, ni ) \
				label( dtrsmi_dtrsm2 )
			cblas_dtrsm( CblasRowMajor,
						 ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Upper,
						 ( CBLAS_TRANSPOSE ) NoTrans,
						 ( CBLAS_DIAG ) NonUnit,
						 tile_size_k, 
						 tile_size_n,
						 1.0, TILE_A[ki][ki], tile_size_k,
							  TILE_B[ki][ni], tile_size_n );
		}

		for ( mi = ki - 1; mi >= 0; mi-- ) 
		{
			tile_size_m = ddss_tile_size( N, mi );
			
			for ( nni = 0; nni < nrhst; nni++ ) 
			{
				tile_size_nn = ddss_tile_size( NRHS, nni );

				#pragma oss task in( TILE_A[mi][ki] ) \
					in( TILE_B[ki][nni] ) \
					inout( TILE_B[mi][nni] ) \
					shared( TILE_A, TILE_B ) \
					firstprivate( ki, mi, nni ) \
					label( dtrsm_dgemm2 )
				cblas_dgemm( CblasRowMajor,
							 CblasNoTrans, CblasNoTrans,
							 tile_size_m, 
							 tile_size_nn, 
							 tile_size_k,
							 -1.0, TILE_A[mi][ki],  tile_size_k,
								   TILE_B[ki][nni], tile_size_nn,
							  1.0, TILE_B[mi][nni], tile_size_nn );
			}
		}
	}

	/*********************************************
	--From tiled data layout to flat data layout--	
	*********************************************/
	ddss_dtiled2flat( N, NRHS, B, LDB, nt, nrhst, TILE_B );

	// --Free--
	free( TILE_A );
	free( TILE_B );
	free( ipiv_local );
	
	return Success;

}
