#include <string.h>

#include <isl/aff.h>

#include "slave_group.h"
#include "slave_array_tile.h"
#include "slave_print.h"
#include "print.h"
#include "schedule.h"

/* Print declarations to "p" for arrays that are local to "prog"
 * but that are used on the host and therefore require a declaration.
 */
__isl_give isl_printer *slave_print_local_declarations(__isl_take isl_printer *p,
  	struct slave_prog *prog)
{
	int i;

	if (!prog)
		return isl_printer_free(p);

	for (i = 0; i < prog->n_array; ++i) {
		struct slave_array_info *array = &prog->array[i];
		isl_ast_expr *size;

		if (!array->declare_local)
			continue;
		size = array->declared_size;
		p = ppcg_print_declaration_with_size(p, array->type, size);
	}

	return p;
}

/* Print an expression for the size of "array" in bytes.
 */
__isl_give isl_printer *slave_array_info_print_size(__isl_take isl_printer *prn,
 	struct slave_array_info *array)
{
	int i;

	for (i = 0; i < array->n_index; ++i) {
		isl_ast_expr *bound;

		prn = isl_printer_print_str(prn, "(");
		bound = isl_ast_expr_get_op_arg(array->bound_expr, 1 + i);
		prn = isl_printer_print_ast_expr(prn, bound);
		isl_ast_expr_free(bound);
		prn = isl_printer_print_str(prn, ") * ");
	}
	prn = isl_printer_print_str(prn, "sizeof(");
	prn = isl_printer_print_str(prn, array->type);
	prn = isl_printer_print_str(prn, ")");

	return prn;
}

__isl_give isl_printer *ppcg_kernel_print_domain_for_sw(__isl_take isl_printer *p,
	struct ppcg_kernel_stmt *stmt)
{
	return pet_stmt_print_body(stmt->u.d.stmt->stmt, p, stmt->u.d.ref2expr);
}

/* Print an access to the element in the private/shared memory copy
 * described by "stmt".  The index of the copy is recorded in
 * stmt->local_index as an access to the array.
 */
__isl_give isl_printer *stmt_print_local_index(__isl_take isl_printer *p,
   	struct ppcg_kernel_stmt *stmt)
{
	return isl_printer_print_ast_expr(p, stmt->u.c.local_index);
}

/* Print an access to the element in the global memory copy
 * described by "stmt".  The index of the copy is recorded in
 * stmt->index as an access to the array.
 */
__isl_give isl_printer *stmt_print_global_index(
  	__isl_take isl_printer *p, struct ppcg_kernel_stmt *stmt)
{
	struct slave_array_info *array = stmt->u.c.array;
	isl_ast_expr *index;

	if (slave_array_is_scalar(array)) {
		if (!slave_array_is_read_only_scalar(array))
			p = isl_printer_print_str(p, "*");
		p = isl_printer_print_str(p, array->name);
		return p;
	}

	index = isl_ast_expr_copy(stmt->u.c.index);

	p = isl_printer_print_ast_expr(p, index);
	isl_ast_expr_free(index);

	return p;
}

/* Assign "aff" to *user and return -1, effectively extracting
 * the first (and presumably only) affine expression in the isl_pw_aff
 * on which this function is used.
 */
static isl_stat extract_single_piece(__isl_take isl_set *set,
	__isl_take isl_aff *aff, void *user)
{
	isl_aff **p = user;

	*p = aff;
	isl_set_free(set);

	return isl_stat_error;
}

/* Print a copy statement.
 *
 * A read copy statement is printed as
 *
 *	local = global;
 *
 * while a write copy statement is printed as
 *
 *	global = local;
 */
/*__isl_give isl_printer *ppcg_kernel_print_copy_for_sw(__isl_take isl_printer *p,
 	struct ppcg_kernel_stmt *stmt)
{
	struct slave_local_array_info *array;
	struct slave_array_ref_group *group;
	struct slave_array_tile *tile;
	isl_val *tsize;
	isl_val *gsize;
	isl_val *bsize;
	int i;
	isl_id *id;
	const char *name;
	const char *type;
	
	isl_pw_aff *pa;
	isl_aff *aff = NULL;

	array = stmt->u.c.local_array;
	group = stmt->u.c.group;
	
	pa = isl_multi_pw_aff_get_pw_aff(group->array->bound, group->array->n_index-1);
	isl_pw_aff_foreach_piece(pa, &extract_single_piece, &aff);
	gsize = isl_aff_get_constant_val(aff);
	tile = group->ldm_tile;
	tsize = isl_val_one(tile->ctx);
	bsize = isl_val_copy(tile->bound[group->array->n_index-1].size);
	for (i = 0; i < tile->n; ++i)
		tsize = isl_val_mul(tsize, isl_val_copy(tile->bound[i].size));
	//tsize = isl_val_copy(tile->bound->size);

	if (stmt->u.c.read) {
		p = isl_printer_start_line(p);
		p = isl_printer_print_str(p, "athread_get(PE_MODE, &");
		p = stmt_print_global_index(p, stmt);
		p = isl_printer_print_str(p, ", ");
		p = isl_printer_print_str(p, "&");
		p = stmt_print_local_index(p, stmt);
		p = isl_printer_print_str(p, ", ");
		p = isl_printer_print_val(p, tsize);
		p = isl_printer_print_str(p, " * ");
		p = isl_printer_print_str(p, "sizeof(");
		p = isl_printer_print_str(p, array->array->type);
		p = isl_printer_print_str(p, "), &");
		p = isl_printer_print_str(p, group->reply->name);
		p = isl_printer_print_str(p, ", 0, ");
		if((isl_val_eq(tsize, bsize) || (isl_val_eq(gsize, bsize)))) {
			p = isl_printer_print_str(p, "0, 0);");
		}
		else {
			p = isl_printer_print_val(p, isl_val_sub(gsize, bsize));
			p = isl_printer_print_str(p, " * ");
			p = isl_printer_print_str(p, "sizeof(");
			p = isl_printer_print_str(p, array->array->type);
			p = isl_printer_print_str(p, "), ");
			p = isl_printer_print_val(p, bsize);
			p = isl_printer_print_str(p, " * ");
			p = isl_printer_print_str(p, "sizeof(");
			p = isl_printer_print_str(p, array->array->type);
			p = isl_printer_print_str(p, "));");
		}
		p = isl_printer_end_line(p);
	} else {
		p = isl_printer_start_line(p);
		p = isl_printer_print_str(p, "athread_put(PE_MODE, &");
		p = stmt_print_local_index(p, stmt);
		p = isl_printer_print_str(p, ", ");
		p = isl_printer_print_str(p, "&");
		p = stmt_print_global_index(p, stmt);
		p = isl_printer_print_str(p, ", ");
		p = isl_printer_print_val(p, tsize);
		p = isl_printer_print_str(p, " * ");
		p = isl_printer_print_str(p, "sizeof(");
		p = isl_printer_print_str(p, array->array->type);
		p = isl_printer_print_str(p, "),&");
		p = isl_printer_print_str(p, group->reply->name);
		p = isl_printer_print_str(p, ", ");
		if((isl_val_eq(tsize, bsize) || (isl_val_eq(gsize, bsize)))) {
			p = isl_printer_print_str(p, "0, 0);");
		}
		else {
			p = isl_printer_print_val(p, isl_val_sub(gsize, bsize));
			p = isl_printer_print_str(p, " * ");
			p = isl_printer_print_str(p, "sizeof(");
			p = isl_printer_print_str(p, array->array->type);
			p = isl_printer_print_str(p, "), ");
			p = isl_printer_print_val(p, bsize);
			p = isl_printer_print_str(p, " * ");
			p = isl_printer_print_str(p, "sizeof(");
			p = isl_printer_print_str(p, array->array->type);
			p = isl_printer_print_str(p, "));");
		}
		p = isl_printer_end_line(p);
	}

	return p;
}*/

/* This function is called for each node in a GPU AST.
 * In case of a user node, print the macro definitions required
 * for printing the AST expressions in the annotation, if any.
 * For other nodes, return true such that descendants are also
 * visited.
 *
 * In particular, for a kernel launch, print the macro definitions
 * needed for the grid size.
 * For a copy statement, print the macro definitions needed
 * for the two index expressions.
 * For an original user statement, print the macro definitions
 * needed for the substitutions.
 */
static isl_bool at_node(__isl_keep isl_ast_node *node, void *user)
{
	const char *name;
	isl_id *id;
	int is_kernel;
	struct ppcg_kernel *kernel;
	struct ppcg_kernel_stmt *stmt;
	isl_printer **p = user;

	if (isl_ast_node_get_type(node) != isl_ast_node_user)
		return isl_bool_true;

	id = isl_ast_node_get_annotation(node);
	if (!id)
		return isl_bool_false;

	name = isl_id_get_name(id);
	if (!name)
		return isl_bool_error;

	if (!strcmp(name, "compute_kernel"))
		return isl_bool_true;

	is_kernel = !strcmp(name, "kernel");
	kernel = is_kernel ? isl_id_get_user(id) : NULL;
	stmt = is_kernel ? NULL : isl_id_get_user(id);
	isl_id_free(id);

	if ((is_kernel && !kernel) || (!is_kernel && !stmt))
		return isl_bool_error;
	
	if (is_kernel) {
		*p = ppcg_ast_expr_print_macros(kernel->kernel_size_expr, *p);
	} else if (stmt->type == ppcg_kernel_copy) {
		*p = ppcg_ast_expr_print_macros(stmt->u.c.index, *p);
		*p = ppcg_ast_expr_print_macros(stmt->u.c.local_index, *p);
	} else if (stmt->type == ppcg_kernel_domain) {
		*p = ppcg_print_body_macros(*p, stmt->u.d.ref2expr);
	}
	if (!*p)
		return isl_bool_error;

	return isl_bool_false;
}

/* Print the required macros for the GPU AST "node" to "p",
 * including those needed for the user statements inside the AST.
 */
__isl_give isl_printer *slave_print_macros(__isl_take isl_printer *p,
 	__isl_keep isl_ast_node *node)
{
	if (isl_ast_node_foreach_descendant_top_down(node, &at_node, &p) < 0)
		return isl_printer_free(p);
	p = ppcg_print_macros(p, node);
	return p;
}						

/* Was the definition of "type" printed before?
 * That is, does its name appear in the list of printed types "types"?
 */
static int already_printed(struct slave_types *types,
 	struct pet_type *type)
{
	int i;

	for (i = 0; i < types->n; ++i)
		if (!strcmp(types->name[i], type->name))
			return 1;

	return 0;
}

/* Print the definitions of all types prog->scop that have not been
 * printed before (according to "types") on "p".
 * Extend the list of printed types "types" with the newly printed types.
 */
__isl_give isl_printer *slave_print_types(__isl_take isl_printer *p,
   	struct slave_types *types, struct slave_prog *prog)
{
	int i, n;
	isl_ctx *ctx;
	char **name;

	n = prog->scop->pet->n_type;

	if (n == 0)
		return p;

	ctx = isl_printer_get_ctx(p);
	name = isl_realloc_array(ctx, types->name, char *, types->n + n);
	if (!name)
		return isl_printer_free(p);
	types->name = name;

	for (i = 0; i < n; ++i) {
		struct pet_type *type = prog->scop->pet->types[i];

		if (already_printed(types, type))
			continue;

		p = isl_printer_start_line(p);
		p = isl_printer_print_str(p, type->definition);
		p = isl_printer_print_str(p, ";");
		p = isl_printer_end_line(p);

		types->name[types->n++] = strdup(type->name);
	}

	return p;
}