#include <core/workers.h>
#include <starpu_sched_node.h>

static double compute_relative_speedup(struct starpu_sched_node * node)
{
	double sum = 0.0;
	int id;
	for(id = starpu_bitmap_first(node->workers_in_ctx);
	    id != -1;
	    id = starpu_bitmap_next(node->workers_in_ctx, id))
	{
		enum starpu_perfmodel_archtype perf_arch = starpu_worker_get_perf_archtype(id);
		sum += starpu_worker_get_relative_speedup(perf_arch);

	}
	STARPU_ASSERT(sum != 0.0);
	return sum;
}


static int push_task(struct starpu_sched_node * node, struct starpu_task * task)
{
	STARPU_ASSERT(node->nchilds > 0);
	int indexes_nodes[node->nchilds];
	double speedup[node->nchilds];
	int size=0,i;
	double alpha_sum = 0.0;
	for(i = 0; i < node->nchilds ; i++)
	{
		if(starpu_sched_node_can_execute_task(node->childs[i],task))
		{
			speedup[size] = compute_relative_speedup(node->childs[i]);
			alpha_sum += speedup[size];
			indexes_nodes[size] = i;
			size++;
		}
	}
	if(size == 0)
		return -ENODEV;

	double random = starpu_drand48()*alpha_sum;
	double alpha = 0.0;
	struct starpu_sched_node * select  = NULL;
	
	for(i = 0; i < size ; i++)
	{
		int index = indexes_nodes[i];
		if(alpha + speedup[i] >= random)
		{	
			select = node->childs[index];
			break;
		}
		alpha += speedup[i];
	}
	STARPU_ASSERT(select != NULL);
	int ret_val = select->push_task(select,task);
	node->available(node);

	return ret_val;
}

double random_estimated_end(struct starpu_sched_node * node)
{
	double sum = 0.0;
	int i;
	for(i = 0; i < node->nchilds; i++)
		sum += node->childs[i]->estimated_end(node->childs[i]);
	return sum / node->nchilds;
}
struct starpu_sched_node * starpu_sched_node_random_create(void * arg STARPU_ATTRIBUTE_UNUSED)
{
	struct starpu_sched_node * node = starpu_sched_node_create();
	node->estimated_end = random_estimated_end;
	node->push_task = push_task;
	return node;
}

int starpu_sched_node_is_random(struct starpu_sched_node *node)
{
	return node->push_task == push_task;
}


static void initialize_random_center_policy(unsigned sched_ctx_id)
{
	starpu_sched_ctx_create_worker_collection(sched_ctx_id, STARPU_WORKER_LIST);
	struct starpu_sched_tree *t = starpu_sched_tree_create(sched_ctx_id);
 	t->root = starpu_sched_node_random_create(NULL);

	unsigned i;
	for(i = 0; i < starpu_worker_get_count() + starpu_combined_worker_get_count(); i++)
	{
		struct starpu_sched_node * node = starpu_sched_node_worker_get(i);
		if(!node)
			continue;
		node->fathers[sched_ctx_id] = t->root;
		t->root->add_child(t->root, node);
	}
	starpu_sched_tree_update_workers(t);
	starpu_sched_ctx_set_policy_data(sched_ctx_id, (void*)t);
}
static void deinitialize_random_center_policy(unsigned sched_ctx_id)
{
	struct starpu_sched_tree *tree = (struct starpu_sched_tree*)starpu_sched_ctx_get_policy_data(sched_ctx_id);
	starpu_sched_tree_destroy(tree, sched_ctx_id);
	starpu_sched_ctx_delete_worker_collection(sched_ctx_id);
}


struct starpu_sched_policy _starpu_sched_tree_random_policy =
{
	.init_sched = initialize_random_center_policy,
	.deinit_sched = deinitialize_random_center_policy,
	.add_workers = starpu_sched_tree_add_workers,
	.remove_workers = starpu_sched_tree_remove_workers,
	.push_task = starpu_sched_tree_push_task,
	.pop_task = starpu_sched_tree_pop_task,
	.pre_exec_hook = NULL,
	.post_exec_hook = NULL,
	.pop_every_task = NULL,
	.policy_name = "tree-random",
	.policy_description = "random tree policy"
};
