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

#include "apl9.h"

Rune primfuncnames[] = L"+-×÷*⍟⌹○!?|⌈⌊⊥⊤⊣⊢=≠≤<>≥≡≢∨∧⍲⍱↑↓⊂⊃⊆⌷⍋⍒⍳⍸∊⍷∪∩~,⍪⍴⌽⊖⍉⍎⍕";

fnmonad monadfunctiondefs[] = {
	0, /* + */
	0, /* - */
	0, /* × */
	0, /* ÷ */
	0, /* * */
	0, /* ⍟ */
	0, /* ⌹ */
	0, /* ○ */
	0, /* ! */
	0, /* ? */
	0, /* | */
	0, /* ⌈ */
	0, /* ⌊ */
	0, /* ⊥ */
	0, /* ⊤ */
	fnSame, /* ⊣ */
	fnSame, /* ⊢ */
	0, /* = */
	0, /* ≠ */
	0, /* ≤ */
	0, /* < */
	0, /* > */
	0, /* ≥ */
	0, /* ≡ */
	fnTally, /* ≢ */
	0, /* ∨ */
	0, /* ∧ */
	0, /* ⍲ */
	0, /* ⍱ */
	0, /* ↑ */
	0, /* ↓ */
	fnEnclose, /* ⊂ */
	0, /* ⊃ */
	fnNest, /* ⊆ */
	0, /* ⌷ */
	0, /* ⍋ */
	0, /* ⍒ */
	fnIndexGenerator, /* ⍳ */
	0, /* ⍸ */
	0, /* ∊ */
	0, /* ⍷ */
	0, /* ∪ */
	0, /* ∩ */
	0, /* ~ */
	fnRavel, /* , */
	fnTable, /* ⍪ */
	fnShape, /* ⍴ */
	fnReverseLast, /* ⌽ */
	fnReverseFirst, /* ⊖ */
	fnTranspose, /* ⍉ */
	0, /* ⍎ */
	0, /* ⍕ */
};

fndyad dyadfunctiondefs[] = {
	fnPlus, /* + */
	fnMinus, /* - */
	fnTimes, /* × */
	fnDivide, /* ÷ */
	fnPower, /* * */
	fnLogarithm, /* ⍟ */
	0, /* ⌹ */
	0, /* ○ */
	0, /* ! */
	0, /* ? */
	0, /* | */
	0, /* ⌈ */
	0, /* ⌊ */
	0, /* ⊥ */
	0, /* ⊤ */
	fnLeft, /* ⊣ */
	fnRight, /* ⊢ */
	0, /* = */
	0, /* ≠ */
	0, /* ≤ */
	0, /* < */
	0, /* > */
	0, /* ≥ */
	fnMatch, /* ≡ */
	0, /* ≢ */
	0, /* ∨ */
	0, /* ∧ */
	0, /* ⍲ */
	0, /* ⍱ */
	0, /* ↑ */
	0, /* ↓ */
	0, /* ⊂ */
	0, /* ⊃ */
	0, /* ⊆ */
	0, /* ⌷ */
	0, /* ⍋ */
	0, /* ⍒ */
	0, /* ⍳ */
	0, /* ⍸ */
	0, /* ∊ */
	0, /* ⍷ */
	0, /* ∪ */
	0, /* ∩ */
	0, /* ~ */
	0, /* , */
	fnCatenateFirst, /* ⍪ */
	fnReshape, /* ⍴ */
	0, /* ⌽ */
	0, /* ⊖ */
	0, /* ⍉ */
	0, /* ⍎ */
	0, /* ⍕ */
};

/* Runner function */
Array *
runfunc(Function f, Array *left, Array *right)
{
	if(f.type == FunctypeDfn){
		Symtab *tmpsymtab = currentsymtab;
		currentsymtab = newsymtab();

		if(left){
			Symbol *alpha = getsym(currentsymtab, L"⍺");
			alpha->value.tag = ArrayTag;
			alpha->value.array = left;
			alpha->undefined = 0;
			incref(left);
		}

		Symbol *omega = getsym(currentsymtab, L"⍵");
		omega->value.tag = ArrayTag;
		omega->value.array = right;
		omega->undefined = 0;
		incref(right);

		Datum *dfnres = evalline(f.dfn);
		freesymtab(currentsymtab);
		currentsymtab = tmpsymtab;
		return (*dfnres).array; /* TODO what if the evaluation failed */
	}else if(f.type == FunctypePrim){
		if(left)
			return dyadfunctiondefs[f.code](left, right);
		else
			return monadfunctiondefs[f.code](right);
	}else{
		/* TODO assumes prim op, not dop */
		if(f.operator.dyadic)
			return dyadoperatordefs[f.operator.code](f.operator.left, f.operator.right, left, right);
		else
			return monadoperatordefs[f.operator.code](f.operator.left, left, right);
	}
}

/* Monadic functions */

Array *
fnSame(Array *right)
{
	incref(right);
	return right;
}

Array *
fnTally(Array *right)
{
	return mkscalarint(right->rank==0 ? 1 : right->shape[0]);
}

Array *
fnEnclose(Array *right)
{
	incref(right);
	if(simplescalar(right))
		return right;
	else{
		Array *res = allocarray(AtypeArray, 0, 1);
		res->arraydata[0] = right;
		return res;
	}
}

Array *
fnIndexGenerator(Array *right)
{
	/* TODO only works for creating vectors */
	vlong n = right->intdata[0];
	Array *res = allocarray(AtypeInt, 1, n);
	res->shape[0] = n;
	vlong io = currentsymtab->io;
	for(vlong i = 0; i < n; i++)
		res->intdata[i] = i + io;
	return res;
}

Array *
fnNest(Array *right)
{
	if(simplearray(right))
		return fnEnclose(right);
	else
		return fnSame(right);
}

Array *
fnRavel(Array *right)
{
	Array *res = duparray(right);
	res->rank = 1;
	res->shape = realloc(res->shape, sizeof(int) * 1);
	res->shape[0] = res->size;
	return res;
}

Array *
fnTable(Array *right)
{
	Array *res = duparray(right);
	res->rank = 2;
	res->shape = realloc(res->shape, sizeof(int) * 2);
	res->shape[0] = right->rank ? res->shape[0] : 1;
	res->shape[1] = right->size / res->shape[0];
	return res;
}

Array *
fnShape(Array *right)
{
	Array *res = allocarray(AtypeInt, 1, right->rank);
	res->shape[0] = right->rank;
	for(int i = 0; i < right->rank; i++)
		res->intdata[i] = right->shape[i];
	return res;
}

Array *
fnReverseLast(Array *right)
{
	if(right->rank < 1)
		return fnSame(right);

	Array *res = duparray(right);

	int nrows = 1;
	int rowsize = res->shape[res->rank-1];
	for(int i = 0; i < res->rank - 1; i++)
		nrows *= res->shape[i];

	for(int row = 0; row < nrows; row++){
		for(int i = 0; i < rowsize; i++)
			memcpy(
				res->rawdata + (row * rowsize + i) * datasizes[res->type],
				right->rawdata + ((1+row) * rowsize - 1 - i) * datasizes[res->type],
				datasizes[res->type]);
	}
	return res;
}

Array *
fnReverseFirst(Array *right)
{
	if(right->rank < 1 || right->shape[0] == 0)
		return fnSame(right);

	Array *res = duparray(right);
	int cells = res->shape[0];
	int elems = res->size / cells;
	for(int i = 0; i < cells; i++)
		memcpy(
			res->rawdata + i * elems * datasizes[res->type],
			right->rawdata + (cells - 1 - i) * elems * datasizes[res->type],
			datasizes[res->type] * elems);
	return res;
}

Array *
fnTranspose(Array *right)
{
	Array *res = duparray(right);
	for(int i = 0; i < res->rank; i++)
		res->shape[i] = right->shape[res->rank - 1 - i];

	int from, to;
	int *sizesFrom = malloc(sizeof(int) * right->rank);
	int *sizesTo = malloc(sizeof(int) * right->rank);
	int accFrom = 1, accTo = 1;
	for(int i = 0; i < right->rank; i++){
		sizesFrom[i] = accFrom;
		sizesTo[i] = accTo;
		accFrom *= res->shape[i];
		accTo *= right->shape[i];
	}

	for(from = 0; from < right->size; from++){
		to = 0;
		int tmp = from;
		for(int i = right->rank-1; i >= 0; i--){
			to += sizesTo[right->rank-1-i] * (tmp / sizesFrom[i]);
			tmp = tmp % sizesFrom[i];
		}
		memcpy(
			res->rawdata + to * datasizes[res->type],
			right->rawdata + from * datasizes[res->type],
			datasizes[res->type]);
	}
	free(sizesFrom);
	free(sizesTo);
	return res;
}

/* Dyadic functions */

Array *
fnPlus(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 0);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++){
		if(res->type == AtypeFloat)
			res->floatdata[i] += right->floatdata[i];
		else if(res->type == AtypeInt)
			res->intdata[i] += right->intdata[i];
	}
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnMinus(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 0);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++){
		if(res->type == AtypeFloat)
			res->floatdata[i] -= right->floatdata[i];
		else if(res->type == AtypeInt)
			res->intdata[i] -= right->intdata[i];
	}
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnTimes(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 0);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++){
		if(res->type == AtypeFloat)
			res->floatdata[i] *= right->floatdata[i];
		else if(res->type == AtypeInt)
			res->intdata[i] *= right->intdata[i];
	}
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnDivide(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 1);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++)
		res->floatdata[i] /= right->floatdata[i];
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnPower(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 0);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++){
		if(res->type == AtypeFloat)
			res->floatdata[i] = pow(res->floatdata[i], right->floatdata[i]);
		else if(res->type == AtypeInt)
			res->intdata[i] = pow(res->intdata[i], right->intdata[i]);
	}
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnLogarithm(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;
	int rankok = scalarextend(left, right, &leftarr, &rightarr);
	if(!rankok){
		print("Ranks don't match lol\n");
		exits(nil);
	}

	int typeok = commontype(leftarr, rightarr, &left, &right, 1);
	if(!typeok){
		print("Types don't match lol\n");
		exits(nil);
	}

	Array *res = duparray(left);
	for(int i = 0; i < left->size; i++)
		res->floatdata[i] = log(right->floatdata[i])/log(res->floatdata[i]);
	freearray(leftarr);
	freearray(rightarr);
	freearray(left);
	freearray(right);
	return res;
}

Array *
fnLeft(Array *left, Array *right)
{
	USED(right);
	incref(left);
	return left;
}

Array *
fnRight(Array *left, Array *right)
{
	USED(left);
	incref(right);
	return right;
}

Array *
fnMatch(Array *left, Array *right)
{
	int cmp = comparearray(left, right);
	return mkscalarint(cmp == 0);
}

Array *
fnCatenateFirst(Array *left, Array *right)
{
	Array *leftarr;
	Array *rightarr;

	if(left->rank == 0 && right->rank != 0){
		/* extend left to right->rank with first axis=1 */
		rightarr = fnSame(right);
		Array *shape = fnShape(right);
		shape->intdata[0] = 1;
		leftarr = fnReshape(shape, left);
		freearray(shape);
	}else if(left->rank != 0 && right->rank == 0){
		/* extend right to left->rank with first axis=1 */
		leftarr = fnSame(left);
		Array *shape = fnShape(left);
		shape->intdata[0] = 1;
		rightarr = fnReshape(shape, right);
		freearray(shape);
	}else if(left->rank == 0 && right->rank == 0){
		/* turn both scalars into vectors */
		leftarr = fnRavel(left);
		rightarr = fnRavel(right);
	}else{
		/* Check that the shapes match */
		if(left->rank == right->rank-1){
			/* extend left with unit dimension */
			Array *shape = allocarray(AtypeInt, 1, left->rank+1);
			shape->intdata[0] = 1;
			for(int i = 1; i < left->rank+1; i++)
				shape->intdata[i] = left->shape[i-1];
			leftarr = fnReshape(shape, left);
			rightarr = fnSame(right);
			freearray(shape);
		}else if(right->rank == left->rank-1){
			/* extend right with unit dimension */
			Array *shape = allocarray(AtypeInt, 1, right->rank+1);
			shape->intdata[0] = 1;
			for(int i = 1; i < right->rank+1; i++)
				shape->intdata[i] = right->shape[i-1];
			rightarr = fnReshape(shape, right);
			leftarr = fnSame(left);
			freearray(shape);
		}else if(right->rank == left->rank){
			leftarr = fnSame(left);
			rightarr = fnSame(right);
		}else{
			print("Ranks don't match\n");
			exits(nil);
			return nil;
		}

		for(int i = 1; i < leftarr->rank; i++)
			if(leftarr->shape[i] != rightarr->shape[i]){
				print("Shapes don't match, lol\n");
				exits(nil);
			}
	}

	int type, rank, leftsize, rightsize;
	if(leftarr->type == AtypeArray || rightarr->type == AtypeArray || leftarr->type != rightarr->type)
		type = AtypeArray;
	else
		type = leftarr->type;
	if(leftarr->rank > rightarr->rank)
		rank = leftarr->rank;
	else
		rank = rightarr->rank;
	
	leftsize = leftarr->shape[0];
	rightsize = rightarr->shape[0];
	for(int i = 1; i < rank; i++)
		leftsize *= leftarr->shape[i];
	for(int i = 1; i < rank; i++)
		rightsize *= rightarr->shape[i];

	Array *result = allocarray(type, rank, leftsize + rightsize);
	int i, j;
	result->shape[0] = leftarr->shape[0] + rightarr->shape[0];
	for(i = 1; i < result->rank; i++)
		result->shape[i] = leftarr->shape[i];

	/* TODO reduce duplicated code between copies from left and right */
	/* Copy data from the left array */
	for(i = 0, j = 0; i < leftarr->size; i++, j++){
		if(type == AtypeArray && leftarr->type == AtypeArray){
			result->arraydata[j] = leftarr->arraydata[i];
			incref(result->arraydata[j]);
		}else if(type == AtypeArray && leftarr->type != AtypeArray){
			result->arraydata[j] = arrayitem(leftarr, i);
		}else{
			memcpy(
				result->rawdata + j * datasizes[type],
				leftarr->rawdata + i * datasizes[type],
				datasizes[type]);
		}
	}

	/* Copy data from the right array */
	for(i = 0; i < rightarr->size; i++, j++){
		if(type == AtypeArray && rightarr->type == AtypeArray){
			result->arraydata[j] = rightarr->arraydata[i];
			incref(result->arraydata[j]);
		}else if(type == AtypeArray && rightarr->type != AtypeArray){
			result->arraydata[j] = arrayitem(rightarr, i);
		}else{
			memcpy(
				result->rawdata + j * datasizes[type],
				rightarr->rawdata + i * datasizes[type],
				datasizes[type]);
		}
	}

	freearray(leftarr);
	freearray(rightarr);
	return result;
}

Array *
fnReshape(Array *left, Array *right)
{
	vlong size = 1;
	int i;
	char *p;
	for(i = 0; i < left->size; i++)
		size *= left->intdata[i];
	if(left->size == 0)
		size = 0;

	Array *res = allocarray(right->type, left->size, size);
	for(i = 0; i < left->size; i++)
		res->shape[i] = left->intdata[i];
	for(i = 0, p = res->rawdata; i < size; i++, p += datasizes[res->type])
		memcpy(p, right->rawdata + (datasizes[res->type] * (i % right->size)), datasizes[res->type]);
	if(res->type == AtypeArray)
		for(i = 0; i < res->size; i++)
			incref(res->arraydata[i]);
	return res;
}