#include <string.h>

#include <isl/schedule.h>
#include <isl/stream.h>
#include <isl_schedule_private.h>
#include <isl_schedule_tree.h>

/* An enumeration of the various keys that may appear in a YAML mapping
 * of a schedule.
 */
enum isl_schedule_key {
	isl_schedule_key_error = -1,
	isl_schedule_key_child,
	isl_schedule_key_coincident,
	isl_schedule_key_context,
	isl_schedule_key_contraction,
	isl_schedule_key_domain,
	isl_schedule_key_expansion,
	isl_schedule_key_extension,
	isl_schedule_key_filter,
	isl_schedule_key_guard,
	isl_schedule_key_leaf,
	isl_schedule_key_mark,
	isl_schedule_key_options,
	isl_schedule_key_permutable,
	isl_schedule_key_schedule,
	isl_schedule_key_sequence,
	isl_schedule_key_set
};

/* Extract a mapping key from the token "tok".
 * Return isl_schedule_key_error on error, i.e., if "tok" does not
 * correspond to any known key.
 */
static enum isl_schedule_key extract_key(__isl_keep isl_stream *s,
	struct isl_token *tok)
{
	int type;
	char *name;
	enum isl_schedule_key key;
	isl_ctx *ctx;

	ctx = isl_stream_get_ctx(s);
	type = isl_token_get_type(tok);
	if (type != ISL_TOKEN_IDENT && type != ISL_TOKEN_STRING) {
		isl_stream_error(s, tok, "expecting key");
		return isl_schedule_key_error;
	}
	name = isl_token_get_str(ctx, tok);
	if (!strcmp(name, "child"))
		key = isl_schedule_key_child;
	else if (!strcmp(name, "coincident"))
		key = isl_schedule_key_coincident;
	else if (!strcmp(name, "context"))
		key = isl_schedule_key_context;
	else if (!strcmp(name, "contraction"))
		key = isl_schedule_key_contraction;
	else if (!strcmp(name, "domain"))
		key = isl_schedule_key_domain;
	else if (!strcmp(name, "expansion"))
		key = isl_schedule_key_expansion;
	else if (!strcmp(name, "extension"))
		key = isl_schedule_key_extension;
	else if (!strcmp(name, "filter"))
		key = isl_schedule_key_filter;
	else if (!strcmp(name, "guard"))
		key = isl_schedule_key_guard;
	else if (!strcmp(name, "leaf"))
		key = isl_schedule_key_leaf;
	else if (!strcmp(name, "mark"))
		key = isl_schedule_key_mark;
	else if (!strcmp(name, "options"))
		key = isl_schedule_key_options;
	else if (!strcmp(name, "schedule"))
		key = isl_schedule_key_schedule;
	else if (!strcmp(name, "sequence"))
		key = isl_schedule_key_sequence;
	else if (!strcmp(name, "set"))
		key = isl_schedule_key_set;
	else if (!strcmp(name, "permutable"))
		key = isl_schedule_key_permutable;
	else
		isl_die(ctx, isl_error_invalid, "unknown key",
			key = isl_schedule_key_error);
	free(name);
	return key;
}

/* Read a key from "s" and return the corresponding enum.
 * Return isl_schedule_key_error on error, i.e., if the first token
 * on the stream does not correspond to any known key.
 */
static enum isl_schedule_key get_key(__isl_keep isl_stream *s)
{
	struct isl_token *tok;
	enum isl_schedule_key key;

	tok = isl_stream_next_token(s);
	key = extract_key(s, tok);
	isl_token_free(tok);

	return key;
}

static __isl_give isl_schedule_tree *isl_stream_read_schedule_tree(
	__isl_keep isl_stream *s);

/* Read a subtree with context root node from "s".
 */
static __isl_give isl_schedule_tree *read_context(__isl_keep isl_stream *s)
{
	isl_set *context = NULL;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	context = isl_set_read_from_str(ctx, str);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		tree = isl_schedule_tree_from_context(context);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_context(tree, context);
	}

	return tree;
error:
	isl_set_free(context);
	return NULL;
}

/* Read a subtree with domain root node from "s".
 */
static __isl_give isl_schedule_tree *read_domain(__isl_keep isl_stream *s)
{
	isl_union_set *domain = NULL;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	domain = isl_union_set_read_from_str(ctx, str);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		tree = isl_schedule_tree_from_domain(domain);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_domain(tree, domain);
	}

	return tree;
error:
	isl_union_set_free(domain);
	return NULL;
}

/* Read a subtree with expansion root node from "s".
 */
static __isl_give isl_schedule_tree *read_expansion(isl_stream *s)
{
	isl_ctx *ctx;
	isl_union_pw_multi_aff *contraction = NULL;
	isl_union_map *expansion = NULL;
	isl_schedule_tree *tree = NULL;
	int more;

	ctx = isl_stream_get_ctx(s);

	do {
		struct isl_token *tok;
		enum isl_schedule_key key;
		char *str;

		key = get_key(s);
		if (isl_stream_yaml_next(s) < 0)
			goto error;

		switch (key) {
		case isl_schedule_key_contraction:
			isl_union_pw_multi_aff_free(contraction);
			tok = isl_stream_next_token(s);
			str = isl_token_get_str(ctx, tok);
			contraction = isl_union_pw_multi_aff_read_from_str(ctx,
									str);
			free(str);
			isl_token_free(tok);
			if (!contraction)
				goto error;
			break;
		case isl_schedule_key_expansion:
			isl_union_map_free(expansion);
			tok = isl_stream_next_token(s);
			str = isl_token_get_str(ctx, tok);
			expansion = isl_union_map_read_from_str(ctx, str);
			free(str);
			isl_token_free(tok);
			if (!expansion)
				goto error;
			break;
		case isl_schedule_key_child:
			isl_schedule_tree_free(tree);
			tree = isl_stream_read_schedule_tree(s);
			if (!tree)
				goto error;
			break;
		default:
			isl_die(ctx, isl_error_invalid, "unexpected key",
				goto error);
		}
	} while ((more = isl_stream_yaml_next(s)) > 0);

	if (more < 0)
		goto error;

	if (!contraction)
		isl_die(ctx, isl_error_invalid, "missing contraction",
			goto error);
	if (!expansion)
		isl_die(ctx, isl_error_invalid, "missing expansion",
			goto error);

	if (!tree)
		return isl_schedule_tree_from_expansion(contraction, expansion);
	return isl_schedule_tree_insert_expansion(tree, contraction, expansion);
error:
	isl_schedule_tree_free(tree);
	isl_union_pw_multi_aff_free(contraction);
	isl_union_map_free(expansion);
	return NULL;
}

/* Read a subtree with extension root node from "s".
 */
static __isl_give isl_schedule_tree *read_extension(isl_stream *s)
{
	isl_union_map *extension = NULL;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	extension = isl_union_map_read_from_str(ctx, str);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		tree = isl_schedule_tree_from_extension(extension);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_extension(tree, extension);
	}

	return tree;
error:
	isl_union_map_free(extension);
	return NULL;
}

/* Read a subtree with filter root node from "s".
 */
static __isl_give isl_schedule_tree *read_filter(__isl_keep isl_stream *s)
{
	isl_union_set *filter = NULL;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	filter = isl_union_set_read_from_str(ctx, str);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		tree = isl_schedule_tree_from_filter(filter);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_filter(tree, filter);
	}

	return tree;
error:
	isl_union_set_free(filter);
	return NULL;
}

/* Read a subtree with guard root node from "s".
 */
static __isl_give isl_schedule_tree *read_guard(isl_stream *s)
{
	isl_set *guard = NULL;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	guard = isl_set_read_from_str(ctx, str);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		tree = isl_schedule_tree_from_guard(guard);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_guard(tree, guard);
	}

	return tree;
error:
	isl_set_free(guard);
	return NULL;
}

/* Read a subtree with mark root node from "s".
 */
static __isl_give isl_schedule_tree *read_mark(isl_stream *s)
{
	isl_id *mark;
	isl_schedule_tree *tree;
	isl_ctx *ctx;
	struct isl_token *tok;
	enum isl_schedule_key key;
	char *str;
	int more;

	ctx = isl_stream_get_ctx(s);

	key = get_key(s);

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	tok = isl_stream_next_token(s);
	if (!tok) {
		isl_stream_error(s, NULL, "unexpected EOF");
		return NULL;
	}
	str = isl_token_get_str(ctx, tok);
	mark = isl_id_alloc(ctx, str, NULL);
	free(str);
	isl_token_free(tok);

	more = isl_stream_yaml_next(s);
	if (more < 0)
		goto error;
	if (!more) {
		isl_die(ctx, isl_error_invalid, "expecting child",
			goto error);
	} else {
		key = get_key(s);
		if (key != isl_schedule_key_child)
			isl_die(ctx, isl_error_invalid, "expecting child",
				goto error);
		if (isl_stream_yaml_next(s) < 0)
			goto error;
		tree = isl_stream_read_schedule_tree(s);
		tree = isl_schedule_tree_insert_mark(tree, mark);
	}

	return tree;
error:
	isl_id_free(mark);
	return NULL;
}

/* Read a sequence of integers from "s" (representing the coincident
 * property of a band node).
 */
static __isl_give isl_val_list *read_coincident(__isl_keep isl_stream *s)
{
	isl_ctx *ctx;
	isl_val_list *list;
	int more;

	ctx = isl_stream_get_ctx(s);

	if (isl_stream_yaml_read_start_sequence(s) < 0)
		return NULL;

	list = isl_val_list_alloc(ctx, 0);
	while  ((more = isl_stream_yaml_next(s)) > 0) {
		isl_val *val;

		val = isl_stream_read_val(s);
		list = isl_val_list_add(list, val);
	}

	if (more < 0 || isl_stream_yaml_read_end_sequence(s))
		list = isl_val_list_free(list);

	return list;
}

/* Set the (initial) coincident properties of "band" according to
 * the (initial) elements of "coincident".
 */
static __isl_give isl_schedule_band *set_coincident(
	__isl_take isl_schedule_band *band, __isl_take isl_val_list *coincident)
{
	int i;
	int n, m;

	n = isl_schedule_band_n_member(band);
	m = isl_val_list_n_val(coincident);

	for (i = 0; i < n && i < m; ++i) {
		isl_val *v;

		v = isl_val_list_get_val(coincident, i);
		if (!v)
			band = isl_schedule_band_free(band);
		band = isl_schedule_band_member_set_coincident(band, i,
							!isl_val_is_zero(v));
		isl_val_free(v);
	}
	isl_val_list_free(coincident);
	return band;
}

/* Read a subtree with band root node from "s".
 */
static __isl_give isl_schedule_tree *read_band(isl_stream *s)
{
	isl_multi_union_pw_aff *schedule = NULL;
	isl_schedule_tree *tree = NULL;
	isl_val_list *coincident = NULL;
	isl_union_set *options = NULL;
	isl_ctx *ctx;
	isl_schedule_band *band;
	int permutable = 0;
	int more;

	ctx = isl_stream_get_ctx(s);

	do {
		struct isl_token *tok;
		enum isl_schedule_key key;
		char *str;
		isl_val *v;

		key = get_key(s);
		if (isl_stream_yaml_next(s) < 0)
			goto error;

		switch (key) {
		case isl_schedule_key_schedule:
			isl_multi_union_pw_aff_free(schedule);
			tok = isl_stream_next_token(s);
			if (!tok) {
				isl_stream_error(s, NULL, "unexpected EOF");
				goto error;
			}
			str = isl_token_get_str(ctx, tok);
			schedule = isl_multi_union_pw_aff_read_from_str(ctx,
									str);
			free(str);
			isl_token_free(tok);
			if (!schedule)
				goto error;
			break;
		case isl_schedule_key_coincident:
			coincident = read_coincident(s);
			if (!coincident)
				goto error;
			break;
		case isl_schedule_key_permutable:
			v = isl_stream_read_val(s);
			permutable = !isl_val_is_zero(v);
			isl_val_free(v);
			break;
		case isl_schedule_key_options:
			isl_union_set_free(options);
			tok = isl_stream_next_token(s);
			str = isl_token_get_str(ctx, tok);
			options = isl_union_set_read_from_str(ctx, str);
			free(str);
			isl_token_free(tok);
			if (!options)
				goto error;
			break;
		case isl_schedule_key_child:
			isl_schedule_tree_free(tree);
			tree = isl_stream_read_schedule_tree(s);
			if (!tree)
				goto error;
			break;
		default:
			isl_die(ctx, isl_error_invalid, "unexpected key",
				goto error);
		}
	} while ((more = isl_stream_yaml_next(s)) > 0);

	if (more < 0)
		goto error;

	if (!schedule)
		isl_die(ctx, isl_error_invalid, "missing schedule", goto error);

	band = isl_schedule_band_from_multi_union_pw_aff(schedule);
	band = isl_schedule_band_set_permutable(band, permutable);
	if (coincident)
		band = set_coincident(band, coincident);
	if (options)
		band = isl_schedule_band_set_ast_build_options(band, options);
	if (tree)
		tree = isl_schedule_tree_insert_band(tree, band);
	else
		tree = isl_schedule_tree_from_band(band);

	return tree;
error:
	isl_val_list_free(coincident);
	isl_union_set_free(options);
	isl_schedule_tree_free(tree);
	isl_multi_union_pw_aff_free(schedule);
	return NULL;
}

/* Read a subtree with root node of type "type" from "s".
 * The node is represented by a sequence of children.
 */
static __isl_give isl_schedule_tree *read_children(isl_stream *s,
	enum isl_schedule_node_type type)
{
	isl_ctx *ctx;
	isl_schedule_tree_list *list;
	int more;

	ctx = isl_stream_get_ctx(s);

	isl_token_free(isl_stream_next_token(s));

	if (isl_stream_yaml_next(s) < 0)
		return NULL;

	if (isl_stream_yaml_read_start_sequence(s))
		return NULL;

	list = isl_schedule_tree_list_alloc(ctx, 0);
	while ((more = isl_stream_yaml_next(s)) > 0) {
		isl_schedule_tree *tree;

		tree = isl_stream_read_schedule_tree(s);
		list = isl_schedule_tree_list_add(list, tree);
	}

	if (more < 0 || isl_stream_yaml_read_end_sequence(s))
		list = isl_schedule_tree_list_free(list);

	return isl_schedule_tree_from_children(type, list);
}

/* Read a subtree with sequence root node from "s".
 */
static __isl_give isl_schedule_tree *read_sequence(isl_stream *s)
{
	return read_children(s, isl_schedule_node_sequence);
}

/* Read a subtree with set root node from "s".
 */
static __isl_give isl_schedule_tree *read_set(isl_stream *s)
{
	return read_children(s, isl_schedule_node_set);
}

/* Read a schedule (sub)tree from "s".
 *
 * We first determine the type of the root node based on the first
 * mapping key and then hand over to a function tailored to reading
 * nodes of this type.
 */
static __isl_give isl_schedule_tree *isl_stream_read_schedule_tree(
	struct isl_stream *s)
{
	enum isl_schedule_key key;
	struct isl_token *tok;
	isl_schedule_tree *tree = NULL;
	int more;

	if (isl_stream_yaml_read_start_mapping(s))
		return NULL;
	more = isl_stream_yaml_next(s);
	if (more < 0)
		return NULL;
	if (!more) {
		isl_stream_error(s, NULL, "missing key");
		return NULL;
	}

	tok = isl_stream_next_token(s);
	key = extract_key(s, tok);
	isl_stream_push_token(s, tok);
	if (key < 0)
		return NULL;
	switch (key) {
	case isl_schedule_key_context:
		tree = read_context(s);
		break;
	case isl_schedule_key_domain:
		tree = read_domain(s);
		break;
	case isl_schedule_key_contraction:
	case isl_schedule_key_expansion:
		tree = read_expansion(s);
		break;
	case isl_schedule_key_extension:
		tree = read_extension(s);
		break;
	case isl_schedule_key_filter:
		tree = read_filter(s);
		break;
	case isl_schedule_key_guard:
		tree = read_guard(s);
		break;
	case isl_schedule_key_leaf:
		isl_token_free(isl_stream_next_token(s));
		tree = isl_schedule_tree_leaf(isl_stream_get_ctx(s));
		break;
	case isl_schedule_key_mark:
		tree = read_mark(s);
		break;
	case isl_schedule_key_sequence:
		tree = read_sequence(s);
		break;
	case isl_schedule_key_set:
		tree = read_set(s);
		break;
	case isl_schedule_key_schedule:
	case isl_schedule_key_coincident:
	case isl_schedule_key_options:
	case isl_schedule_key_permutable:
		tree = read_band(s);
		break;
	case isl_schedule_key_child:
		isl_die(isl_stream_get_ctx(s), isl_error_unsupported,
			"cannot identity node type", return NULL);
	case isl_schedule_key_error:
		return NULL;
	}

	if (isl_stream_yaml_read_end_mapping(s) < 0) {
		isl_stream_error(s, NULL, "unexpected extra elements");
		return isl_schedule_tree_free(tree);
	}

	return tree;
}

/* Read an isl_schedule from "s".
 */
__isl_give isl_schedule *isl_stream_read_schedule(isl_stream *s)
{
	isl_ctx *ctx;
	isl_schedule_tree *tree;

	if (!s)
		return NULL;

	ctx = isl_stream_get_ctx(s);
	tree = isl_stream_read_schedule_tree(s);
	return isl_schedule_from_schedule_tree(ctx, tree);
}

/* Read an isl_schedule from "input".
 */
__isl_give isl_schedule *isl_schedule_read_from_file(isl_ctx *ctx, FILE *input)
{
	struct isl_stream *s;
	isl_schedule *schedule;

	s = isl_stream_new_file(ctx, input);
	if (!s)
		return NULL;
	schedule = isl_stream_read_schedule(s);
	isl_stream_free(s);

	return schedule;
}

/* Read an isl_schedule from "str".
 */
__isl_give isl_schedule *isl_schedule_read_from_str(isl_ctx *ctx,
	const char *str)
{
	struct isl_stream *s;
	isl_schedule *schedule;

	s = isl_stream_new_str(ctx, str);
	if (!s)
		return NULL;
	schedule = isl_stream_read_schedule(s);
	isl_stream_free(s);

	return schedule;
}