#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),
	[AtypeMixed] = sizeof(Mixed),
	[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 = duparrayshape(a, a->type);
	memcpy(b->rawdata, a->rawdata, datasizes[a->type]*a->size);
	if(b->type == AtypeArray)
		for(int i = 0; i < b->size; i++)
			incarrayref(b->arraydata[i]);
	return b;
}

Array *
duparrayshape(Array *a, int type)
{
	Array *b = allocarray(type, a->rank, a->size);
	memcpy(b->shape, a->shape, sizeof(int) * a->rank);
	b->stranded = a->stranded;
	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->rank == 0 && b->rank != 0){
		*aa = extend(a, b);
		*bb = fnSame(b);
	}else if(a->rank != 0 && b->rank == 0){
		*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 if(a->size == 1 && b->size == 1){
		Array *shape;
		if(a->rank > b->rank)
			shape = fnShape(a);
		else
			shape = fnShape(b);
		*aa = fnReshape(shape, a);
		*bb = fnReshape(shape, b);
		freearray(shape);
	}else if(a->size == 1 && b->size != 1){
		Array *shape = fnShape(b);
		*aa = fnReshape(shape, a);
		*bb = fnSame(b);
		freearray(shape);
	}else if(a->size != 1 && b->size == 1){
		Array *shape = fnShape(a);
		*aa = fnSame(a);
		*bb = fnReshape(shape, b);
		freearray(shape); 
	}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 AtypeMixed:
		switch(a->mixeddata[index].type){
		case AtypeInt: res = mkscalarint(a->mixeddata[index].i); break;
		case AtypeFloat: res = mkscalarfloat(a->mixeddata[index].f); break;
		case AtypeRune: res = mkscalarrune(a->mixeddata[index].r); break;
		default:
			throwerror(L"Unhandled case in arrayitem", ENotImplemented);
		}
		break;
	case AtypeArray:
		res = a->arraydata[index];
		incarrayref(res);
		break;
	default:
		throwerror(L"Unhandled case in arrayitem", ENotImplemented);
	}
	return res;
}

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

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

	if(sametype && type != AtypeArray){
		Array *b = duparrayshape(a, type);
		for(i = 0; i < a->size; i++){
			if(nested){
				memcpy(b->rawdata + i * datasizes[type], a->arraydata[i]->rawdata, datasizes[type]);
				if(b->type == AtypeArray)
					incarrayref(b->arraydata[i]);
			}else{
				switch(b->type){
				case AtypeInt: b->intdata[i] = a->mixeddata[i].i; break;
				case AtypeFloat: b->floatdata[i] = a->mixeddata[i].f; break;
				case AtypeRune: b->runedata[i] = a->mixeddata[i].r; break;
				default:
					throwerror(L"Missing case in simplifyarray", ENotImplemented);
				}
			}
		}
		if(b->type == AtypeArray){
			Array *tmp = b;
			b = simplifyarray(b);
			freearray(tmp);
		}
		return b;
	}else if(canfloat){
		Array *b = duparrayshape(a, AtypeFloat);
		for(i = 0; i < a->size; i++){
			if(nested){
				if(a->arraydata[i]->type == AtypeFloat)
					b->floatdata[i] = a->arraydata[i]->floatdata[0];
				else
					b->floatdata[i] = a->arraydata[i]->intdata[0];
			}else{
				if(a->mixeddata[i].type == AtypeFloat)
					b->floatdata[i] = a->mixeddata[i].f;
				else
					b->floatdata[i] = a->mixeddata[i].i;
			}
		}
		return b;
	}else if(canmix && nested){
		Array *b = duparrayshape(a, AtypeMixed);
		for(i = 0; i < a->size; i++){
			b->mixeddata[i].type = a->arraydata[i]->type;
			switch(b->mixeddata[i].type){
			case AtypeInt: b->mixeddata[i].i = a->arraydata[i]->intdata[0]; break;
			case AtypeFloat: b->mixeddata[i].f = a->arraydata[i]->floatdata[0]; break;
			case AtypeRune: b->mixeddata[i].r = a->arraydata[i]->runedata[0]; break;
			default:
				throwerror(L"Missing case in simplifyarray", ENotImplemented);
			}
		}
		return b;
	}else
		return fnSame(a);
}

int
comparearray(Array *a, Array *b, int checkshapes)
{
	/* returns -1 if a < b, 0 if a == b and 1 if a > b. */
	int i;

	if(a->type < b->type)
		return -1;
	else if(a->type > b->type)
		return 1;

	if(checkshapes){
		if(a->rank < b->rank)
			return -1;
		else if(a->rank > b->rank)
			return 1;
		for(i = 0; i < a->rank; i++){
			if(a->shape[i] < b->shape[i])
				return -1;
			else if(a->shape[i] > b->shape[i])
				return 1;
		}
	}

	for(i = 0; i < a->size && i < b->size; i++){
		int sub = 0;
		switch(a->type){
		case AtypeInt:
			sub = a->intdata[i] > b->intdata[i] ? 1 : a->intdata[i] == b->intdata[i] ? 0 : -1;
			break;
		case AtypeFloat:
			sub = a->floatdata[i] > b->floatdata[i] ? 1 : a->floatdata[i] == b->floatdata[i] ? 0 : -1;
			break;
		case AtypeRune:
			sub = a->runedata[i] > b->runedata[i] ? 1 : a->runedata[i] == b->runedata[i] ? 0 : -1;
			break;
		case AtypeArray:
			sub = comparearray(a->arraydata[i], b->arraydata[i], checkshapes);
			break;
		default:
			print("Missing comparison code for type %d\n", a->type);
			exits(nil);
		}
		if(sub != 0)
			return sub;
	}

	if(i < a->size)
		return 1;
	else if(i < b->size)
		return -1;
	else
		return 0;
}

Array *
fillelement(Array *a)
{
	switch(a->type){
	case AtypeInt: return mkscalarint(0);
	case AtypeFloat: return mkscalarfloat(0);
	case AtypeRune: return mkscalarrune(' ');
	case AtypeMixed:{
		Array *first = arrayitem(a, 0);
		Array *fill = fillelement(first);
		freearray(first);
		return fill;
	}
	case AtypeArray:{
		Array *b = duparrayshape(a, a->type);
		for(int i = 0; i < b->size; i++){
			Array *fill = fillelement(a->arraydata[i]);
			Array *shape = fnShape(a->arraydata[i]);
			b->arraydata[i] = fnReshape(shape, fill);
			freearray(fill);
			freearray(shape);
		}
		return b;
	}
	default:
		print("Can't make fill element of array type %d\n", a->type);
		exits(nil);
		return 0;
	}
}

uvlong
arrayspaceused(Array *a)
{
	uvlong size = 0;
	size += sizeof(*a);
	size += sizeof(int) * a->rank;
	size += datasizes[a->type] * a->size;

	for(int i = 0; i < a->size && a->type == AtypeArray; i++)
		size += arrayspaceused(a->arraydata[i]);
	return size;
}

int
arraydepth(Array *a, int *uniform)
{
	if(a->type == AtypeArray){
		int max = -1;
		int subuniform;
		*uniform = 1;
		for(int i = 0; i < a->size; i++){
			int subdepth = arraydepth(a->arraydata[i], &subuniform);
			if((subdepth != max && max != -1) || subuniform == 0)
				*uniform = 0;
			if(subdepth > max)
				max = subdepth;
		}
		return max+1;
	}else{
		*uniform = 1;
		if(a->rank == 0)
			return 0;
		else
			return 1;
	}
}