/**
 *This is the source code for Mars, a MapReduce framework on graphics
 *processors.
 *Author: Wenbin Fang (HKUST), Bingsheng He (HKUST)
 *Mentor: Naga K. Govindaraju (Microsoft Corp.), Qiong Luo (HKUST), Tuyong
 *Wang (Sina.com).
 *If you have any question on the code, please contact us at {saven,
 *wenbin, luo}@cse.ust.hk.
 *The copyright is held by HKUST. Mars is provided "as is" without any 
 *guarantees of any kind.
 */
 
#undef __CPU_MAP__
#undef __GPU_MAP__
#undef __CPU_REDUCE__
#undef __GPU_REDUCE__
#define __COMPARE__
#include "MarsInc.h"
#include "MarsInc/MarsConfig.h"
 
//==============================================================
//CPU sort
//==============================================================
ChunkInfo_t *g_chunk = NULL;

int cmp_wrap(const void *arg1, const void *arg2)
{
	int4 *arg_1 = (int4*)arg1;
	int4 *arg_2 = (int4*)arg2;

	char *key1 = g_chunk->keys + arg_1->x + g_chunk->keyOffset;
	char *key2 = g_chunk->keys + arg_2->x + g_chunk->keyOffset;
	size_t key1Size = arg_1->y;
	size_t key2Size = arg_2->y;
	int ret = cpu_compare(key1, key1Size, key2, key2Size);
	return ret;
}

void InitGroupMem(ChunkInfo_t *chunk)
{
	BEN_ASSERT(chunk != NULL);
	g_chunk = chunk;
}
 
void QuickSortMem(ChunkInfo_t *chunk)
{
	BEN_ASSERT(chunk != NULL);
	g_chunk = chunk;
	qsort(chunk->index, chunk->recCount,sizeof(int4), 
		cmp_wrap);
}
 
void GroupByMem(ChunkInfo_t *chunk)
{
	BEN_ASSERT(chunk != NULL);
		
	chunk->keyListRange = (int2*)BenMalloc(sizeof(int2)*chunk->recCount);
	chunk->diffKeyCount = 1;
	chunk->keyListRange[0].x = 0;
	int4 *preOffsetSize = chunk->index;
 
	int i;
	for (i = 1; i < chunk->recCount; i++)
	{
		if (cmp_wrap(preOffsetSize, &chunk->index[i]) != 0)
		{
			chunk->keyListRange[chunk->diffKeyCount].x = i;
			chunk->keyListRange[chunk->diffKeyCount-1].y = i;
			chunk->diffKeyCount++;
			preOffsetSize = &(chunk->index[i]);
		}
	}
	chunk->keyListRange[chunk->diffKeyCount-1].y = i;
	chunk->keyListRange = (int2*)BenRealloc(chunk->keyListRange, sizeof(int2)*chunk->recCount, sizeof(int2)*chunk->diffKeyCount);
}

//--------------------------------------------------------------
//external merge sort
//--------------------------------------------------------------
static char GetSmallSortChunk(FileName_t *file,
							  SortChunk_t *smallChunk, 
							  SortChunk_t *bigChunk, 
							  size_t chunkSize,
							  char *bitmap, int index,
							  size_t *recOffsets)
{
	BEN_ASSERT(smallChunk != NULL);
	BEN_ASSERT(bigChunk != NULL);
	BEN_ASSERT(file != NULL);
	BEN_ASSERT(bitmap != NULL);
	BEN_ASSERT(recOffsets != NULL);

	if (smallChunk[index].cursor >= bigChunk[index].diffKeyCount)
	{
		bitmap[index] = 0;
		return 0;
	}

	smallChunk[index].diffKeyCount = chunkSize;
	if (smallChunk[index].diffKeyCount + smallChunk[index].cursor >
		bigChunk[index].diffKeyCount)
		smallChunk[index].diffKeyCount = bigChunk[index].diffKeyCount - 
		smallChunk[index].cursor;

	//read rangeFile
	smallChunk[index].rangeOffset = sizeof(int2)*(smallChunk[index].cursor)+bigChunk[index].rangeOffset;
	smallChunk[index].rangeSize = sizeof(int2)*smallChunk[index].diffKeyCount;
	smallChunk[index].keyListRange = (int2*)BenReadFile(file->rangeFile, 
		smallChunk[index].rangeOffset, 
		smallChunk[index].rangeSize);
	smallChunk[index].rangeOffset = sizeof(int2)*smallChunk[index].cursor;

	//read indexFile
	smallChunk[index].recCount = smallChunk[index].keyListRange[smallChunk[index].diffKeyCount-1].y
		- smallChunk[index].keyListRange[0].x;
	smallChunk[index].indexOffset = sizeof(int4)*smallChunk[index].keyListRange[0].x+bigChunk[index].indexOffset;
	smallChunk[index].indexSize = sizeof(int4)*smallChunk[index].recCount;
	smallChunk[index].index = (int4*)BenReadFile(file->indexFile, 
		smallChunk[index].indexOffset,
		smallChunk[index].indexSize);
	smallChunk[index].indexOffset = sizeof(int4)*smallChunk[index].keyListRange[0].x;
	
	//read keyFile
	smallChunk[index].keyOffset = smallChunk[index].index[0].x+bigChunk[index].keyOffset;
	smallChunk[index].keySize = smallChunk[index].index[smallChunk[index].recCount - 1].x +
		smallChunk[index].index[smallChunk[index].recCount - 1].y;
	smallChunk[index].keys = (char*)BenReadFile(file->keyFile, 
		smallChunk[index].keyOffset,
		smallChunk[index].keySize);
	smallChunk[index].keyOffset = smallChunk[index].index[0].x;

	//read valFile
	smallChunk[index].valOffset = smallChunk[index].index[0].z+bigChunk[index].valOffset;
	smallChunk[index].valSize = smallChunk[index].index[smallChunk[index].recCount - 1].z +
		smallChunk[index].index[smallChunk[index].recCount - 1].w;
	smallChunk[index].vals = (char*)BenReadFile(file->valFile, 
		smallChunk[index].valOffset,
		smallChunk[index].valSize);
	smallChunk[index].valOffset = smallChunk[index].index[0].z;

	smallChunk[index].cursor += smallChunk[index].diffKeyCount;

	recOffsets[index] += smallChunk[index].recCount;

	bitmap[index] = 1;

	return 1;
}

//used only when ((GPU_SORT or GPU_SORT_REDUCE) and USE_FILE)
void RearrangeKeyVal(ChunkInfo_t *chunk)
{
	char *newKeys = (char*)BenMalloc(chunk->keySize);
	char *newVals = (char*)BenMalloc(chunk->valSize);

	size_t keyOffset = 0;
	size_t valOffset = 0;
	for (int i = 0; i < chunk->recCount; i++)
	{
		int4 index = chunk->index[i];
		//BenLog("%d,%d\n", *(int*)(chunk->keys+index.x),
		//	*(int*)(chunk->vals+index.z));
		BenMemcpy(newKeys+keyOffset, chunk->keys+index.x, index.y);
		BenMemcpy(newVals+valOffset, chunk->vals+index.z, index.w);
		chunk->index[i].x = keyOffset;
		chunk->index[i].z = valOffset;
		keyOffset += index.y;
		valOffset += index.w;
	}

	BenFree((char**)&chunk->keys, chunk->keySize);
	BenFree((char**)&chunk->vals, chunk->valSize);
	chunk->keys = newKeys;
	chunk->vals = newVals;
}

//get a group of records from small sort chunk
static void GetSortRec(SortRec_t *rec, SortChunk_t *chunks, 
					   int index, size_t *recOffsets)
{
	BEN_ASSERT(rec != NULL);
	BEN_ASSERT(chunks != NULL);

	SortChunk_t *chunk = &chunks[index];
	if (chunk->smallCursor >= chunk->diffKeyCount)
		return;
	size_t i = chunk->smallCursor;

	rec->keyListRange = chunk->keyListRange[i];
	rec->valCount = rec->keyListRange.y - rec->keyListRange.x;
	rec->index = &chunk->index[rec->keyListRange.x-recOffsets[index]+ chunk->recCount];
	rec->allValSize = 0;
	for (int j = 0; j < rec->valCount; j++)
		rec->allValSize += rec->index[j].w;
	rec->keySize = rec->index[0].y;
	rec->allKeySize = rec->valCount * rec->keySize;
	rec->key = chunk->keys + rec->index[0].x - chunk->keyOffset ;
	rec->val = chunk->vals + rec->index[0].z - chunk->valOffset ;
}

//--------------------------------------------------------------
//add record, group, 
//--------------------------------------------------------------
//static offset control
static int4 outputSize = make_int4(DEFAULT_SORT_OUTPUT_SIZE,
	DEFAULT_SORT_OUTPUT_SIZE, DEFAULT_SORT_OUTPUT_SIZE,
	DEFAULT_SORT_OUTPUT_SIZE);
//0,1,2,3,4,5, ... 0, 1,2, 3,...0,1,2,...
static int4 perRecSizeOffset = make_int4(0,0,0,0);
//0,0,0,0,1,1,1,1,2,2,2,2, 333,444
static int2 perChunkSizeOffset = make_int2(0,0);
//1,2,3,4,5...n for index
static size_t recOffset = 0;
//0,0,0, 1,1,1,...
static size_t curOffset = 0;
static size_t groupOffset = 0;
//0, 1, 2, 3, ... n
static int2 keyValOffset = make_int2(0, 0);
//used in WriteChunk
static size_t totalRecCount = 0;
static size_t totalGroupCount = 0;
static void OutputAndGroupRecs(ChunkInfo_t *outBuf,
							   FileName_t *file,
							   SortRec_t *rec, 
							   size_t chunkSize)
{
	BEN_ASSERT(outBuf != NULL);
	BEN_ASSERT(chunkSize > 0);
	BEN_ASSERT(rec != NULL);
	BEN_ASSERT(file != NULL);

	//flag symbol control
	static char firstRec = 1;
	static char firstGroup = 1;

	//allocate buffer
	if (firstRec == 1)
	{
		outBuf->keys = (char*)BenMalloc(outputSize.x);
		outBuf->vals = (char*)BenMalloc(outputSize.y);
		outBuf->index = (int4*)BenMalloc(outputSize.z);
		outBuf->keyListRange = (int2*)BenMalloc(outputSize.w);
		
		firstRec = 0;
	}

	//key
	if (perRecSizeOffset.x + rec->allKeySize >= outputSize.x)
	{
		outBuf->keys = (char*)BenRealloc(outBuf->keys, outputSize.x, outputSize.x*2);
		outputSize.x *= 2;
	}
	BenMemcpy(outBuf->keys+perRecSizeOffset.x, rec->key, rec->allKeySize);

	//val
	if (perRecSizeOffset.y + rec->allValSize >= outputSize.y)
	{
		outBuf->vals = (char*)BenRealloc(outBuf->vals, outputSize.y, outputSize.y*2);
		outputSize.y *= 2;
	}
	BenMemcpy(outBuf->vals+perRecSizeOffset.y, rec->val, rec->allValSize);

	//index
	size_t indexSize = rec->valCount*sizeof(int4);
	if (perRecSizeOffset.z + indexSize >= outputSize.z)
	{
		outBuf->index = (int4*)BenRealloc(outBuf->index, outputSize.z, outputSize.z*2);
		outputSize.z *= 2;
	}
	for (int i = 0; i < rec->valCount; i++)
	{
		rec->index[i].x = keyValOffset.x;
		rec->index[i].z = keyValOffset.y;
		keyValOffset.x += rec->index[i].y;
		keyValOffset.y += rec->index[i].w;
	}
	BenMemcpy(((char*)outBuf->index)+perRecSizeOffset.z, rec->index, indexSize);

	//keyListRange
	if (perRecSizeOffset.w + sizeof(int2) >= outputSize.w)
	{
		outBuf->keyListRange = (int2*)BenRealloc(outBuf->keyListRange, outputSize.w, outputSize.w*2);
		outputSize.w *= 2;
	}
		
	//group
	static int4 preIndex = *rec->index;
	int4 curIndex = *rec->index;
	if (preIndex.x >= perChunkSizeOffset.x)
	{
		preIndex.x -= perChunkSizeOffset.x;
		preIndex.z -= perChunkSizeOffset.y;
	}
	if (curIndex.x >= perChunkSizeOffset.x)
	{
		curIndex.x -= perChunkSizeOffset.x;
		curIndex.z -= perChunkSizeOffset.y;
	}

	//!!!
	if (firstGroup != 1)
	{
		InitGroupMem(outBuf);
		size_t backKeyOffset = outBuf->keyOffset;
		size_t backValOffset = outBuf->valOffset;
		outBuf->keyOffset = 0;
		outBuf->valOffset = 0;

		if (cmp_wrap(&preIndex, &curIndex) != 0)
		{
			outBuf->keyListRange[outBuf->diffKeyCount].x = recOffset;
			outBuf->keyListRange[outBuf->diffKeyCount-1].y = recOffset;
			totalGroupCount ++;			
			preIndex = outBuf->index[recOffset - curOffset];
			outBuf->diffKeyCount++;
			perRecSizeOffset.w += sizeof(int2);
			outBuf->rangeSize += sizeof(int2);
		}

		outBuf->keyOffset = backKeyOffset;
		outBuf->valOffset = backValOffset;
	}
	else
	{
		firstGroup = 0;
		perRecSizeOffset.w += sizeof(int2);	
		outBuf->diffKeyCount = 1;
		outBuf->keyListRange[0].x = recOffset;
		outBuf->rangeSize += sizeof(int2);
	}
	//!!!

	//proceed
	//adjust chunk
	outBuf->keySize += rec->allKeySize;
	outBuf->valSize += rec->allValSize;
	outBuf->indexSize += indexSize;
	outBuf->recCount += rec->valCount;

	//adjust offset
	perRecSizeOffset.x += rec->allKeySize;
	perRecSizeOffset.y += rec->allValSize;
	perRecSizeOffset.z += indexSize;
	recOffset += rec->valCount;
	groupOffset++;

	//flush
	if (outBuf->diffKeyCount >= chunkSize)
	{
		//flush
		char mode = EXTERNAL_SORT;
		outBuf->keySize -= rec->allKeySize;
		outBuf->valSize -= rec->allValSize;
		outBuf->indexSize -= rec->valCount*sizeof(int4);
		outBuf->rangeSize -= sizeof(int2);
		outBuf->recCount -= rec->valCount;

		WriteChunkToFile(outBuf, file, &totalRecCount, mode);

		//adjust offset
		perRecSizeOffset = make_int4(rec->allKeySize,
			rec->allValSize, rec->valCount*sizeof(int4), sizeof(int2));
		perChunkSizeOffset.x += outBuf->keySize;
		perChunkSizeOffset.y += outBuf->valSize;

		curOffset += outBuf->recCount;
		recOffset = outBuf->keyListRange[outBuf->diffKeyCount-1].x;
		preIndex = outBuf->index[0];
		
		//adjust chunk
		BenMemcpy(outBuf->keys, rec->key, rec->allKeySize);
		BenMemcpy(outBuf->vals, rec->val, rec->allValSize);
		BenMemcpy(outBuf->index, rec->index, rec->valCount*sizeof(int4));
		
		outBuf->recCount = rec->valCount;
		outBuf->keySize = rec->allKeySize;
		outBuf->valSize = rec->allValSize;
		outBuf->indexSize = rec->valCount*sizeof(int4);
		outBuf->rangeSize = sizeof(int2);
		outBuf->diffKeyCount = 1;
		outBuf->keyListRange[0].x = recOffset;
		recOffset += rec->valCount;
	}
}

//------------------------------------------------------------
//the last flush of sort chunk
//------------------------------------------------------------
static void FlushSortChunk(ChunkInfo_t *chunk, FileName_t *file)
{
	BEN_ASSERT(chunk != NULL);
	BEN_ASSERT(file != NULL);

	char mode = EXTERNAL_SORT;
	chunk->keyListRange[chunk->diffKeyCount-1].y = recOffset;
	WriteChunkToFile(chunk, file, &totalRecCount, mode);
	totalGroupCount ++;

	BenFree((char**)&chunk->keys, outputSize.x);
	BenFree((char**)&chunk->vals, outputSize.y);
	BenFree((char**)&chunk->index, outputSize.z);
	BenFree((char**)&chunk->keyListRange, outputSize.w);
}

void MergeSortFile(FileName_t *file, 
				   FileName_t *tmpfile,
				   SortInfo_t *sortInfo, 
				   size_t chunkRecCount,
				   size_t *totalOutputRecCount,
				   size_t *totalOutputDiffKeyCount)
{
	//EnterFunc("MergeSortFile");

	BEN_ASSERT(file != NULL);
	BEN_ASSERT(sortInfo != NULL);
	BEN_ASSERT(chunkRecCount > 0);

	SortChunk_t *smallChunks = 
		(SortChunk_t*)BenMalloc
		(sizeof(SortChunk_t)*sortInfo->realChunkCount);
	char *bigChunkBitmap = (char*)BenMalloc(sortInfo->realChunkCount);