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

#include "apl9.h"

Rune primmonopnames[] = L"¨⍨⌸⌶&⌾∆⍇";
Rune primdyadopnames[] = L"⍣.∘⍤⍥@⍠⌺⍢⍫⍙";

opmonad monadoperatordefs[] = {
	opEach, /* ¨ */
	opSwitch, /* ⍨ */
	opKey, /* ⌸ */
	0, /* ⌶ */
	opSpawn, /* & */
	opOuterProduct, /* ⌾ */
	opSelfReference1, /* ∆ */
	opReceive, /* ⍇ */
};

opdyad dyadoperatordefs[] = {
	opPower, /* ⍣ */
	opInnerProduct, /* . */
	opBind, /* ∘ */
	opAtop, /* ⍤ */
	opOver, /* ⍥ */
	0, /* @ */
	0, /* ⍠ */
	0, /* ⌺ */
	opUnder, /* ⍢ */
	opObverse, /* ⍫ */
	opSelfReference2, /* ⍙ */
};

/* Monadic operators */
Array *
opEach(Datum *lefto, Array *left, Array *right)
{
	int i;
	if(lefto->tag != FunctionTag)
		return nil;
	Array *leftarr;
	Array *rightarr;
	if(left){
		int rankok = scalarextend(left, right, &leftarr, &rightarr);
		if(!rankok)
			throwerror(nil, ERank);
	}else{
		leftarr = nil;
		rightarr = fnSame(right);
	}

	/* TODO handle empty arrays by applying the function to their prototype */

	Array *result = allocarray(AtypeArray, GetRank(rightarr), GetSize(rightarr));
	for(i = 0; i < GetRank(rightarr); i++)
		result->shape[i] = rightarr->shape[i];
	for(i = 0; i < GetSize(rightarr); i++){
		Array *elem1 = leftarr ? arrayitem(leftarr, i) : nil;
		Array *elem2 = arrayitem(rightarr, i);
		result->arraydata[i] = runfunc(lefto->func, elem1, elem2);
		freearray(elem1);
		freearray(elem2);
	}
	freearray(leftarr);
	freearray(rightarr);
	return result;
}

Array *
opSwitch(Datum *lefto, Array *left, Array *right)
{
	if(lefto->tag == ArrayTag){
		incarrayref(lefto->array);
		return lefto->array;
	}else if(lefto->tag == FunctionTag){
		if(left)
			return runfunc(lefto->func, right, left);
		else
			return runfunc(lefto->func, right, right);
	}else
		return nil;
}

Array *
opKey(Datum *lefto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag)
		throwerror(nil, EType);
	if(left)
		return rundfn(L"↑⍵∘(⍶{⍵⍶⍺⌷⍨⊂⍸⍹≡¨⍵}⍺)¨⊂¨∪⍺", lefto, nil, left, right);
	else
		return rundfn(L"⍵⍶⌸⍳≢⍵", lefto, nil, left, right);
}

Array *
opSpawn(Datum *lefto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag)
		throwerror(L"Can only spawn functions", EType);
	int id = spawnthread(lefto->func, left, right);
	return mkscalarint(id);
}

Array *
opOuterProduct(Datum *lefto, Array *left, Array *right)
{
	if(left == nil)
		throwerror(L"f⌾ expected a left argument", ESyntax);
	if(lefto->tag != FunctionTag)
		throwerror(nil, EType);

	int i;
	int rank = GetRank(left) + GetRank(right);
	int size = 1;
	int *shape = emalloc(sizeof(int) * rank);
	for(i = 0; i < GetRank(left); i++){
		shape[i] = left->shape[i];
		size *= left->shape[i];
	}
	for(i = 0; i < GetRank(right); i++){
		shape[i+GetRank(left)] = right->shape[i];
		size *= right->shape[i];
	}
	Array *result = allocarray(AtypeArray, rank, size);
	for(i = 0; i < rank; i++)
		result->shape[i] = shape[i];

	for(int leftindex = 0; leftindex < GetSize(left); leftindex++){
		for(int rightindex = 0; rightindex < GetSize(right); rightindex++){
			i = leftindex * GetSize(right) + rightindex;
			Array *leftitem = arrayitem(left, leftindex);
			Array *rightitem = arrayitem(right, rightindex);
			result->arraydata[i] = runfunc(lefto->func, leftitem, rightitem);
			freearray(leftitem);
			freearray(rightitem);
		}
	}
	free(shape);
	return result;
}

Array *
opSelfReference1(Datum *lefto, Array *left, Array *right)
{
	DfnFrame *dfn = getcurrentdfn();
	if(dfn)
		return rundfn(dfn->code, lefto, nil, left, right);
	else{
		throwerror(nil, ESyntax);
		return nil;
	}
}

Array *
opReceive(Datum *lefto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag)
		throwerror(nil, EType);
	if(GetType(right) != AtypeInt)
		throwerror(nil, EType);
	if(GetSize(right) != 1)
		throwerror(nil, ELength);
	USED(left);
	return messagerecv(lefto->func, right->intdata[0]);
}

/* Dyadic operators */
Array *
opPower(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	Rune *code = nil;
	if(lefto->tag != FunctionTag)
		throwerror(nil, EType);
	if(righto->tag == FunctionTag){
		if(left)
			code = L"next←⍺⍶⍵ ⋄ next⍹⍵:⍵ ⋄ ⍺∇next";
		else
			code = L"next←⍶⍵ ⋄ next⍹⍵:⍵ ⋄ ∇next";
	}else if(righto->tag == ArrayTag){
		if(GetType(righto->array) != AtypeInt || GetRank(righto->array) != 0 || GetSize(righto->array) != 1)
			throwerror(L"right operand to ⍣", EDomain);
		vlong times = righto->array->intdata[0];
		if(times < 0){
			lefto->func = inverse(lefto->func);
			times = -times;
		}
		righto->array->intdata[0] = times;
		if(left)
			code = L"⍹=0:⍵ ⋄ ⍺ ⍶⍙(⍹-1)⊢⍺⍶⍵";
		else
			code = L"⍹=0:⍵ ⋄ ⍶⍙(⍹-1)⊢⍶⍵";
	}

	if(code)
		return rundfn(code, lefto, righto, left, right);
	else{
		throwerror(nil, EDomain);
		return nil;
	}
}

Array *
opInnerProduct(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(left == nil)
		throwerror(L"f.g expected a left argument", ESyntax);
	if(lefto->tag != FunctionTag || righto->tag != FunctionTag)
		throwerror(nil, EType);

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

	return rundfn(L"↑(↓⍺)(⍶/⍹¨)⌾⍉↓⍉⍵", lefto, righto, left, right);
}

Array *
opBind(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(lefto->tag == FunctionTag && righto->tag == FunctionTag){
		Array *right1 = runfunc(righto->func, nil, right);
		Array *result = runfunc(lefto->func, left, right1);
		freearray(right1);
		return result;
	}else if(lefto->tag == FunctionTag && righto->tag == ArrayTag){
		if(left)
			throwerror(L"The function doesn't take a left argument", ESyntax);
		return runfunc(lefto->func, right, righto->array);
	}else if(lefto->tag == ArrayTag && righto->tag == FunctionTag){
		if(left)
			throwerror(L"The function doesn't take a left argument", ESyntax);
		return runfunc(righto->func, lefto->array, right);
	}else
		return nil;
}

Array *
opAtop(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(lefto->tag == FunctionTag && righto->tag == FunctionTag){
		Array *right1;
		if(left)
			right1 = runfunc(righto->func, left, right);
		else
			right1 = runfunc(righto->func, nil, right);
		Array *result = runfunc(lefto->func, nil, right1);
		freearray(right1);
		return result;
	}else if(lefto->tag == FunctionTag && righto->tag == ArrayTag){
		Array *ranks = righto->array;
		if(GetRank(ranks) > 1)
			throwerror(nil, ERank);
		if(GetType(ranks) != AtypeInt)
			throwerror(nil, EType);
		if(GetSize(ranks) < 1 || GetSize(ranks) > 3)
			throwerror(nil, ELength);
		int p,q,r;
		switch(GetSize(ranks)){
		case 1: p = q = r = ranks->intdata[0]; break;
		case 2:
			q = ranks->intdata[0];
			p = r = ranks->intdata[1];
			break;
		case 3:
			p = ranks->intdata[0];
			q = ranks->intdata[1];
			r = ranks->intdata[2];
			break;
		default:
			p = q = r = 0;
		}

		/* "Normalize" the ranks (remove negatives, and ranks above the array's ranks */
		if(left){
			q = q < 0 ? GetRank(left) + q : q;
			q = q < 0 ? 0 : q;
			q = q > GetRank(left) ? GetRank(left) : q;
		}
		p = p < 0 ? GetRank(right) + p : p;
		p = p < 0 ? 0 : p;
		p = p > GetRank(right) ? GetRank(right) : p;
		r = r < 0 ? GetRank(right) + r : r;
		r = r < 0 ? 0 : r;
		r = r > GetRank(right) ? GetRank(right) : r;

		Array *result;
		if(left){
			Array *cellrankl = mkscalarint(q);
			Array *cellrankr = mkscalarint(r);
			Array *lefts = rundfn(L"(⌷∘⍵)¨⍳(-⍺)↓⍴⍵", nil, nil, cellrankl, left);
			Array *rights = rundfn(L"(⌷∘⍵)¨⍳(-⍺)↓⍴⍵", nil, nil, cellrankr, right);
			result = rundfn(L"↑⍺⍶¨⍵", lefto, nil, lefts, rights);
			freearray(cellrankl);
			freearray(cellrankr);
			freearray(lefts);
			freearray(rights);
		}else{
			if(GetRank(right) == p)
				result = runfunc(lefto->func, nil, right);
			else{
				Array *cellrank = mkscalarint(p);
				result = rundfn(L"↑⍶¨(⌷∘⍵)¨⍳(-⍺)↓⍴⍵", lefto, nil, cellrank, right);
				freearray(cellrank);
			}
		}
		return result;
	}else{
		throwerror(L"Unknown combination of operands to ⍤", ENotImplemented);
		return nil;
	}
}

Array *
opUnder(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag || righto->tag != FunctionTag)
		throwerror(L"⍢ operands must be functions", EDomain);
	if(left)
		return rundfn(L"⍹⍣¯1 ⊢ (⍹ ⍺) ⍶ ⍹ ⍵", lefto, righto, left, right);
	else
		return rundfn(L"⍹⍣¯1 ⍶ ⍹ ⍵", lefto, righto, left, right);
}

Array *
opObverse(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag || righto->tag != FunctionTag)
		throwerror(L"⍫ operands must be functions", EDomain);
	return runfunc(lefto->func, left, right);
}

Array *
opOver(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	if(lefto->tag != FunctionTag || righto->tag != FunctionTag)
		return nil;
	if(left){
		Array *r = runfunc(righto->func, nil, right);
		Array *l = runfunc(righto->func, nil, left);
		Array *res = runfunc(lefto->func, l, r);
		freearray(r);
		freearray(l);
		return res;
	}else{
		Array *tmp = runfunc(righto->func, nil, right);
		Array *res = runfunc(lefto->func, nil, tmp);
		freearray(tmp);
		return res;
	}
}

Array *
opSelfReference2(Datum *lefto, Datum *righto, Array *left, Array *right)
{
	DfnFrame *dfn = getcurrentdfn();
	if(dfn)
		return rundfn(dfn->code, lefto, righto, left, right);
	else{
		throwerror(nil, ESyntax);
		return nil;
	}
}