#include <u.h>
#include <libc.h>
#include <bio.h>

#include "apl9.h"

Rune primhybridnames[] = L"/\\⌿⍀";

fndyad hybridfunctiondefs[] = {
	fnReplicateLast, /* / */
	fnExpandLast, /* \ */
	fnReplicateFirst, /* ⌿ */
	fnExpandFirst, /* ⍀ */
};

opmonad hybridoperatordefs[] = {
	opReduceLast, /* / */
	opScanLast, /* \ */
	opReduceFirst, /* ⌿ */
	opScanFirst, /* ⍀ */
};

/* function definitions */
Array *
fnReplicateLast(Array *left, Array *right)
{
	return rundfn(L"⍉⍺⌿⍉⍵", nil, nil, left, right);
}

Array *
fnExpandLast(Array *left, Array *right)
{
	return rundfn(L"⍉⍺⍀⍉⍵", nil, nil, left, right);
}

Array *
fnReplicateFirst(Array *left, Array *right)
{
	if(left->type != AtypeInt)
		throwerror(nil, EType);
	if(left->rank > 1)
		throwerror(nil, ERank);

	/* Reshape right if scalar */
	if(right->size == 1){
		Array *shape = right->rank == 0 ? mkscalarint(1) : fnShape(right);
		shape->intdata[0] = left->size;
		right = fnReshape(shape, right);
		freearray(shape);
	}else
		right = fnSame(right);

	int nsize = right->shape[0];
	int cellsize = nsize == 0 ? 1 : right->size / nsize;
	int i;
	if(left->size == 1){
		Array *shape = mkscalarint(nsize);
		left = fnReshape(shape, left);
	}else
		left = fnSame(left);

	if(left->size != nsize){
		freearray(left);
		freearray(right);
		throwerror(nil, ELength);
	}

	nsize = 0;
	for(i = 0; i < left->size; i++){
		vlong n = left->intdata[i];
		nsize += n > 0 ? n : -n;
	}

	Array *result = allocarray(right->type, right->rank, nsize * cellsize);
	result->shape[0] = nsize;
	for(i = 1; i < result->rank; i++)
		result->shape[i] = right->shape[i];

	Array *fill = fillelement(right);
	int from, to;
	int npos;
	for(from = 0, to = 0; to < nsize; from++, to += npos){
		int neg = left->intdata[from] < 0;
		npos = left->intdata[from];
		if(neg){
			npos = -npos;
			for(int j = 0; j < npos*cellsize; j++)
				memcpy(result->rawdata + (to*cellsize+j)*datasizes[result->type],
					fill->rawdata, datasizes[result->type]);
		}else{
			for(int j = 0; j < npos; j++)
				memcpy(result->rawdata + (to+j)*cellsize*datasizes[result->type],
					right->rawdata + from*cellsize*datasizes[result->type],
					cellsize*datasizes[result->type]);
		}
		if(result->type == AtypeArray)
			for(int j = 0; j < npos*cellsize; j++)
				incref(result->arraydata[to*cellsize+j]);
	}
	freearray(fill);
	freearray(left);
	freearray(right);
	return result;
}

Array *
fnExpandFirst(Array *left, Array *right)
{
	if(left->type != AtypeInt)
		throwerror(nil, EType);
	if(left->rank > 1)
		throwerror(nil, ERank);

	right = right->rank == 0 ? fnRavel(right) : fnSame(right);

	int npos = 0;
	int nsize = 0;
	int i;
	for(i = 0; i < left->size; i++){
		if(left->intdata[i] > 0)
			npos++;
		else if(left->intdata[i] == 0)
			left->intdata[i] = -1;

		vlong n = left->intdata[i];
		nsize += n > 0 ? n : -n;
	}

	if(right->shape[0] != 1 && right->shape[0] != npos)
		throwerror(nil, ELength);

	vlong cellsize = right->shape[0] == 0 ? 1 : right->size / right->shape[0];
	Array *result = allocarray(right->type, right->rank, nsize * cellsize);
	result->shape[0] = nsize;
	for(i = 1; i < result->rank; i++)
		result->shape[i] = right->shape[i];

	Array *fill = fillelement(right);
	int from, to;
	for(from = 0, to = 0, i = 0; to < nsize; i++, to += npos){
		int neg = left->intdata[i] < 0;
		npos = left->intdata[i];
		if(neg){
			npos = -npos;
			for(int j = 0; j < npos*cellsize; j++)
				memcpy(result->rawdata + (to*cellsize+j)*datasizes[result->type],
					fill->rawdata, datasizes[result->type]);
		}else{
			for(int j = 0; j < npos; j++)
				memcpy(result->rawdata + (to+j)*cellsize*datasizes[result->type],
					right->rawdata + from*cellsize*datasizes[result->type],
					cellsize*datasizes[result->type]);
		}
		if(result->type == AtypeArray)
			for(int j = 0; j < npos*cellsize; j++)
				incref(result->arraydata[to*cellsize+j]);
		if(right->shape[0] != 1 && !neg)
			from++;
	}
	freearray(fill);
	freearray(right);
	return result;
}

/* operator definitions */
Array *
opReduceLast(Datum *lefto, Array *left, Array *right)
{
	if(left)
		throwerror(L"left argument to f/", ENotImplemented);

	right = fnTranspose(right);
	Array *tmp = opReduceFirst(lefto, left, right);
	Array *res = fnTranspose(tmp);
	freearray(right);
	freearray(tmp);
	return res;
}

Array *
opScanLast(Datum *lefto, Array *left, Array *right)
{
	if(left)
		throwerror(L"f\ doesn't take a left argument", ESyntax);

	right = fnTranspose(right);
	Array *tmp = opScanFirst(lefto, left, right);
	Array *res = fnTranspose(tmp);
	freearray(right);
	freearray(tmp);
	return res;
}

Array *
opReduceFirst(Datum *lefto, Array *left, Array *right)
{
	if(left)
		throwerror(L"left argument to f⌿", ENotImplemented);

	if(right->rank == 0)
		return fnSame(right);

	int n = right->shape[0];
	int io = globalIO();
	if(n == 0)
		throwerror(L"Can't figure out identity element", ENotImplemented);

	Array *result = allocarray(AtypeArray, right->rank - 1, right->size / n);
	for(int i = 0; i < right->rank-1; i++)
		result->shape[i] = right->shape[i+1];

	Array *index = allocarray(AtypeArray, 1, right->rank);
	index->shape[0] = right->rank;
	Array *tmp = mkscalarint(n);
	index->arraydata[0] = fnIndexGenerator(tmp);
	freearray(tmp);
	for(int i = 1; i < index->size; i++)
		index->arraydata[i] = mkscalarint(io);

	for(int i = 0; i < result->size; i++){
		for(int j = index->size - 1; index->arraydata[j]->intdata[0] == io + right->shape[j]; j--){
			index->arraydata[j]->intdata[0] = io;
			index->arraydata[j-1]->intdata[0]++;
		}
		Array *vector = fnIndex(index, right);
		result->arraydata[i] = arrayitem(vector, n-1);
		for(int j = n-2; j >= 0; j--){
			Array *argL = arrayitem(vector, j);
			Array *argR = result->arraydata[i];
			result->arraydata[i] = runfunc(lefto->func, argL, argR);
			freearray(argL);
			freearray(argR);
		}
		index->arraydata[index->size-1]->intdata[0]++;
	}
	freearray(index);

	return result;
}

Array *
opScanFirst(Datum *lefto, Array *left, Array *right)
{
	if(left)
		throwerror(L"f⍀ doesn't take a left argument", ESyntax);

	Array *result = allocarray(AtypeArray, right->rank, right->size);
	for(int i = 0; i < right->rank; i++)
		result->shape[i] = right->shape[i];

	int n = result->shape[0];
	int m = result->size / n;
	for(int i = 0; i < n; i++){
		Array *len = mkscalarint(i + 1);
		Array *index = fnIndexGenerator(len);
		Array *ix = fnEnclose(index);
		Array *subarr = fnIndex(ix, right);
		Array *subres = opReduceFirst(lefto, left, subarr);
		for(int j = 0; j < m; j++)
			result->arraydata[i*m + j] = arrayitem(subres, j);
		freearray(len);
		freearray(index);
		freearray(ix);
		freearray(subarr);
		freearray(subres);
	}
	return result;
}