// Copyright (c) 2023, Marvin Borner <dev@marvinborner.de>
// SPDX-License-Identifier: MIT

#include <math.h>
#include <time.h>
#include <stdlib.h>
#include <assert.h>

#include <schedule.h>
#include <reduce.h>
#include <lib/pqueue.h>
#include <log.h>
#include <parse.h>
#include <map.h>

// queue of beta-redexes: ([X], Y)
static struct pqueue *queue;

static pqueue_pri_t get_pri(void *a)
{
	struct term *term = a;
	return term->priority;
}

static void set_pri(void *a, pqueue_pri_t pri)
{
	struct term *term = a;
	term->priority = pri;
}

static int cmp_pri(pqueue_pri_t next, pqueue_pri_t curr)
{
	return next < curr;
}

static size_t get_pos(void *a)
{
	struct term *term = a;
	return term->position;
}

static void set_pos(void *a, size_t position)
{
	struct term *term = a;
	term->position = position;
}

// skewed random number generator
// prefers smaller numbers
static size_t choose_position(void)
{
	size_t size = pqueue_size(queue);
	return size - (size_t)sqrt(((size_t)rand() % (size_t)pow(size, 2)));
}

static pqueue_pri_t calculate_priority(struct term *term)
{
	// TODO: Try different formulas (hyperfine)
	return (parse_get_max_depth() - term->depth + 1) * term->refs;
}

void schedule_add(struct term *term)
{
	if (!term_is_beta_redex(term))
		return;

	set_pri(term, calculate_priority(term));
	pqueue_insert(queue, term);
}

void schedule_remove(struct term *term)
{
	if (!pqueue_size(queue) || !term_is_beta_redex(term))
		return;

	int error = pqueue_remove(queue, term);
	assert(!error);
}

void schedule(void)
{
	while (pqueue_size(queue)) {
		debug("queue size: %zu\n", pqueue_size(queue));
		map_dump(map_all_terms(), 1);

		// TODO: check finished programs
		size_t position = choose_position();
		struct term *term = pqueue_pop_at(queue, position);
		reduce(term);
	}
	debug("no more redexes!\n");
}

void schedule_sync_priorities(void)
{
	size_t iter = 0;
	void *iter_val;
	while (hashmap_iter(map_all_terms(), &iter, &iter_val)) {
		struct term *term = *(struct term **)iter_val;
		if (!term_is_beta_redex(term))
			continue;
		pqueue_change_priority(queue, calculate_priority(term), term);
	}
}

void schedule_initialize(void)
{
	srand(time(0));
	size_t size = hashmap_count(map_all_terms()) + 42;
	queue = pqueue_init(size, cmp_pri, get_pri, set_pri, get_pos, set_pos);
}

void schedule_destroy(void)
{
	pqueue_free(queue);
}