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

#include "apl9.h"

Array *inttofloatarray(Array *);

int datasizes[] = {
	[AtypeInt] = sizeof(vlong),
	[AtypeFloat] = sizeof(double),
	[AtypeRune] = sizeof(Rune),
	[AtypeArray] = sizeof(Array *)
};

Array *
mkscalarint(vlong i)
{
	Array *a = allocarray(AtypeInt, 0, 1);
	a->intdata[0] = i;
	return a;
}

Array *
mkscalarfloat(double f)
{
	Array *a = allocarray(AtypeFloat, 0, 1);
	a->floatdata[0] = f;
	return a;
}

Array *
mkscalarrune(Rune r)
{
	Array *a = allocarray(AtypeRune, 0, 1);
	a->runedata[0] = r;
	return a;
}

Array *
mkrunearray(Rune *str)
{
	Array *a = allocarray(AtypeRune, 1, runestrlen(str));
	a->shape[0] = a->size;
	for(int i = 0; i < a->size; i++)
		a->runedata[i] = str[i];
	return a;
}

Array *
duparray(Array *a)
{
	Array *b = allocarray(a->type, a->rank, a->size);
	memcpy(b->shape, a->shape, sizeof(int) * a->rank);
	memcpy(b->rawdata, a->rawdata, datasizes[a->type]*a->size);
	if(b->type == AtypeArray)
		for(int i = 0; i < b->size; i++)
			incref(b->arraydata[i]);
	return b;
}

int
simplearray(Array *a)
{
	return a->type != AtypeArray;
}

int
simplescalar(Array *a)
{
	return a->rank == 0 && a->type != AtypeArray;
}

Array *
extend(Array *a, Array *b)
{
	/* extend the singleton a to the shape of b */
	Array *shape = fnShape(b);
	Array *res = fnReshape(shape, a);
	freearray(shape);
	return res;
}

int
scalarextend(Array *a, Array *b, Array **aa, Array **bb)
{
	/* Extend the arrays a and b to have the same shape.
	   The resulting arrays are stored in aa and bb,
	   except when the ranks don't match or extension can't
	   happen, in which case the function returns 0 and
	   aa and bb are unchanged.
	*/

	if(a->size == 1 && b->size != 1){
		*aa = extend(a, b);
		*bb = fnSame(b);
	}else if(b->size == 1 && a->size != 1){
		*aa = fnSame(a);
		*bb = extend(b, a);
	}else if(a->size == b->size && a->rank == b->rank){
		/* Check that each dimension matches */
		for(int i = 0; i < a->rank; i++)
			if(a->shape[i] != b->shape[i])
				return 0;
		*aa = fnSame(a);
		*bb = fnSame(b);
	}else
		return 0;
	return 1;
}

Array *
inttofloatarray(Array *a)
{
	Array *b = allocarray(AtypeFloat, a->rank, a->size);
	for(int i = 0; i < a->rank; i++)
		b->shape[i] = a->shape[i];
	for(int i = 0; i < a->size; i++)
		b->floatdata[i] = a->intdata[i];
	return b;
}

int
commontype(Array *a, Array *b, Array **aa, Array **bb, int forcefloat)
{
	/* When A and B are numeric arrays, set aa and bb
	   to arrays that have compatible types, with the same data.
	*/
	if(forcefloat){
		*aa = a->type == AtypeFloat ? fnSame(a) : inttofloatarray(a);
		*bb = b->type == AtypeFloat ? fnSame(b) : inttofloatarray(b);
	}else if(a->type == b->type){
		*aa = fnSame(a);
		*bb = fnSame(b);
	}else if(a->type == AtypeFloat && b->type == AtypeInt){
		*aa = fnSame(a);
		*bb = inttofloatarray(b);
	}else if(a->type == AtypeInt && b->type == AtypeFloat){
		*aa = inttofloatarray(a);
		*bb = fnSame(b);
	}else
		return 0;
	return 1;
}

Array *
arrayitem(Array *a, int index)
{
	Array *res = nil;
	switch(a->type){
	case AtypeInt:
		res = mkscalarint(a->intdata[index]);
		break;
	case AtypeFloat:
		res = mkscalarfloat(a->floatdata[index]);
		break;
	case AtypeRune:
		res = mkscalarrune(a->runedata[index]);
		break;
	case AtypeArray:
		res = a->arraydata[index];
		incref(res);
		break;
	default:
		print("Unhandled case in arrayitem()\n");
		exits(nil);
	}
	return res;
}

Array *
simplifyarray(Array *a)
{
	/* simplify an array if possible. */
	if(a->type != AtypeArray || a->size == 0)
		return fnSame(a);
	int type = a->arraydata[0]->type;
	int canfloat = type == AtypeFloat || type == AtypeInt;
	int sametype = 1;
	int i;

	for(i = 0; i < a->size; i++){
		int t = a->arraydata[i]->type;
		canfloat = canfloat && (t == AtypeFloat || t == AtypeInt);
		sametype = sametype && (t == type);
		if(a->arraydata[i]->rank != 0)
			return fnSame(a); /* cannot be simplified */
	}

	if(sametype){
		Array *b = allocarray(type, a->rank, a->size);
		b->stranded = a->stranded;
		for(i = 0; i < a->rank; i++)
			b->shape[i] = a->shape[i];
		for(i = 0; i < a->size; i++)
			memcpy(b->rawdata + i * datasizes[type], a->arraydata[i]->rawdata, datasizes[type]);
		return b;
	}else if(canfloat){
		Array *b = allocarray(AtypeFloat, a->rank, a->size);
		b->stranded = a->stranded;
		for(i = 0; i < a->rank; i++)
			b->shape[i] = a->shape[i];
		for(i = 0; i < a->size; i++){
			if(a->arraydata[i]->type == AtypeFloat)
				b->floatdata[i] = a->arraydata[i]->floatdata[0];
			else
				b->floatdata[i] = a->arraydata[i]->intdata[0];
		}
		return b;
	}else
		return fnSame(a);
}

int
comparearray(Array *a, Array *b)
{
	/* returns -1 if a < b, 0 if a == b and 1 if a > b.
	   Only correctly handles test for equality right now,
	   and returns 1 for unequal data. */

	if(a->type != b->type)
		return 1;
	if(a->rank != b->rank)
		return 1;
	if(a->size != b->size)
		return 1;
	for(int i = 0; i < a->rank; i++)
		if(a->shape[i] != b->shape[i])
			return 1;
	for(int i = 0; i < a->size; i++){
		if(a->type == AtypeArray){
			/* do something recursive here */
			int sub = comparearray(a->arraydata[i], b->arraydata[i]);
			if(sub != 0)
				return sub;
		}else{
			int sub = memcmp(
					a->rawdata + i * datasizes[a->type],
					b->rawdata + i * datasizes[a->type],
					datasizes[a->type]);
			if(sub < 0)
				return -1;
			else if(sub > 0)
				return 1;	
		}
	}

	/* if we get here, the arrays are equal */
	return 0;
}