#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(GetType(left) != AtypeInt)
		throwerror(nil, EType);
	if(GetRank(left) > 1)
		throwerror(nil, ERank);

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

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

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

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

	Array *result = allocarray(GetType(right), GetRank(right), nsize * cellsize);
	result->shape[0] = nsize;
	for(i = 1; i < GetRank(result); 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[GetType(result)],
					fill->rawdata, datasizes[GetType(result)]);
		}else{
			for(int j = 0; j < npos; j++)
				memcpy(result->rawdata + (to+j)*cellsize*datasizes[GetType(result)],
					right->rawdata + from*cellsize*datasizes[GetType(result)],
					cellsize*datasizes[GetType(result)]);
		}
		if(GetType(result) == AtypeArray)
			for(int j = 0; j < npos*cellsize; j++)
				incarrayref(result->arraydata[to*cellsize+j]);
	}
	freearray(fill);
	freearray(left);
	freearray(right);
	return result;
}

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

	right = GetRank(right) == 0 ? fnRavel(right) : fnSame(right);

	int npos = 0;
	int nsize = 0;
	int i;
	for(i = 0; i < GetSize(left); 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 : GetSize(right) / right->shape[0];
	Array *result = allocarray(GetType(right), GetRank(right), nsize * cellsize);
	result->shape[0] = nsize;
	for(i = 1; i < GetRank(result); 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[GetType(result)],
					fill->rawdata, datasizes[GetType(result)]);
		}else{
			for(int j = 0; j < npos; j++)
				memcpy(result->rawdata + (to+j)*cellsize*datasizes[GetType(result)],
					right->rawdata + from*cellsize*datasizes[GetType(result)],
					cellsize*datasizes[GetType(result)]);
		}
		if(GetType(result) == AtypeArray)
			for(int j = 0; j < npos*cellsize; j++)
				incarrayref(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)
{
	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){
		if(GetType(left) != AtypeInt)
			throwerror(nil, EType);
		if(GetSize(left) != 1)
			throwerror(nil, ELength);
		vlong winsize = left->intdata[0];
		if(winsize > right->shape[0])
			throwerror(nil, EShape);

		Rune *code = L"n←(-|⍺)+1+≢⍵ ⋄"
				L"ix←(⍳|⍺)∘+¨(⍳n)-⎕io ⋄"
				L"ix←⍺{⍺<0:⌽¨⍵ ⋄ ⍵}ix ⋄"
				L"↑⍵∘(⍶{⍶⌿(⊂⍵)⌷⍺})¨ix";
		return rundfn(code, lefto, nil, left, right);
	}

	if(GetRank(right) == 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, GetRank(right) - 1, GetSize(right) / n);
	for(int i = 0; i < GetRank(right)-1; i++)
		result->shape[i] = right->shape[i+1];

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

	for(int i = 0; i < GetSize(result); i++){
		for(int j = GetSize(index) - 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);
		}
		freearray(vector);
		index->arraydata[GetSize(index)-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 = duparrayshape(right, AtypeArray);
	int n = result->shape[0];
	int m = GetSize(result) / 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;
}