#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, /* , */
	0, /* ⍪ */
	fnShape, /* ⍴ */
	0, /* ⌽ */
	0, /* ⊖ */
	0, /* ⍉ */
	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, /* ≥ */
	0, /* ≡ */
	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, /* ⍕ */
};

/* 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 = globalIO();
	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 *
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;
}

/* 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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] += rightarr->intdata[i];
	freearray(leftarr);
	freearray(rightarr);
	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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] -= rightarr->intdata[i];
	freearray(leftarr);
	freearray(rightarr);
	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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] *= rightarr->intdata[i];
	freearray(leftarr);
	freearray(rightarr);
	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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] /= rightarr->intdata[i];
	freearray(leftarr);
	freearray(rightarr);
	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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] = pow(res->intdata[i], rightarr->intdata[i]);
	freearray(leftarr);
	freearray(rightarr);
	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);
	}

	Array *res = duparray(leftarr);
	for(int i = 0; i < leftarr->size; i++)
		res->intdata[i] = log(rightarr->intdata[i])/log(res->intdata[i]);
	freearray(leftarr);
	freearray(rightarr);
	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 *
fnCatenateFirst(Array *left, Array *right)
{
	/* not even close to being right, but it works for stranding :) */
	left = left->rank == 0 ? fnRavel(left) : fnSame(left);
	right = right->rank == 0 ? fnRavel(right) : fnSame(right);

	/* assume two vectors of same type for now */
	Array *res = allocarray(left->type, 1, left->size+right->size);
	res->shape[0] = left->shape[0] + right->shape[0];
	memcpy(res->rawdata, left->rawdata, datasizes[res->type]*left->size);
	memcpy(res->rawdata+datasizes[res->type]*left->size, right->rawdata, datasizes[res->type]*right->size);
	if(res->type == AtypeArray)
		for(int i = 0; i < res->size; i++)
			incref(res->arraydata[i]);
	freearray(left);
	freearray(right);
	return res;
}

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;
}