#include "online_crossbar.h"

#include "crossbar_generator.h"
#include "read_crossbar.h"

#include <stdbool.h>

#include <ctype.h>
#include <math.h>
#include <ngspice/sharedspice.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct Online_Crossbar {
	Online_Crossbar_Config config;
	double *resistances;
	conductance_mapping mapping;
	double *state_ring;
	double *results;
	size_t published_count;
	size_t consumed_count;
	size_t result_count;
	int started;
	int simulation_done;
	int failed;
	int mutex_initialized;
	int state_condition_initialized;
	int progress_condition_initialized;
	pthread_mutex_t mutex;
	pthread_cond_t state_available;
	pthread_cond_t progress;
};

// checks that the online crossbar configuration is valid
static int config_is_valid(const Online_Crossbar_Config *config)
{
	if (!config || config->num_neurons == 0 || config->num_outputs == 0 ||
	    config->num_timesteps == 0 || !isfinite(config->time_step) ||
	    config->time_step <= 0.0 || !isfinite(config->spike_amplitude) ||
	    config->spike_amplitude == 0.0 ||
	    !isfinite(config->load_resistance) ||
	    config->load_resistance <= 0.0 || !isfinite(config->r_on) ||
	    !isfinite(config->r_off) || config->r_on <= 0.0 ||
	    config->r_off <= config->r_on || !config->model_path ||
	    !config->subcircuit_name || !config->netlist_path) {
		return 0;
	}

	if (config->num_outputs > SIZE_MAX / 2 ||
	    config->num_neurons > SIZE_MAX / 2 ||
	    config->num_timesteps > SIZE_MAX / config->num_outputs ||
	    config->num_timesteps * config->num_outputs >
		SIZE_MAX / sizeof(double)) {
		return 0;
	}

	double stop_time = config->num_timesteps * config->time_step;
	return isfinite(stop_time) && stop_time > 0.0;
}

// ignores random ngspice output, unless its an error
static int callback_text(char *text, int ident, void *user_data)
{
	(void)ident;
	(void)user_data;
	if (text &&
	    (strncmp(text, "stderr", 6) == 0 || strstr(text, "Error") != NULL ||
	     strstr(text, "error") != NULL)) {
		fprintf(stderr, "ngspice: %s\n", text);
	}
	return 0;
}

// ignores ngspice vector information that the benchmark does not need
static int callback_init_data(pvecinfoall values, int ident, void *user_data)
{
	(void)values;
	(void)ident;
	(void)user_data;
	return 0;
}

// marks the simulation as failed and wakes any waiting threads
static void fail_locked(Online_Crossbar *crossbar)
{
	crossbar->failed = 1;
	pthread_cond_broadcast(&crossbar->state_available);
	pthread_cond_broadcast(&crossbar->progress);
}

// marks the simulation as failed if ngspice exits unexpectedly
static int callback_exit(int status, NG_BOOL immediate, NG_BOOL quit_exit,
			 int ident, void *user_data)
{
	(void)status;
	(void)immediate;
	(void)quit_exit;
	(void)ident;
	Online_Crossbar *crossbar = user_data;
	pthread_mutex_lock(&crossbar->mutex);
	fail_locked(crossbar);
	pthread_mutex_unlock(&crossbar->mutex);
	return 0;
}

// records when the ngspice background simulation finishes
static int callback_background(NG_BOOL running, int ident, void *user_data)
{
	(void)ident;
	Online_Crossbar *crossbar = user_data;
	/* libngspice passes false at worker start and true at worker exit. */
	if (!running) {
		return 0;
	}
	pthread_mutex_lock(&crossbar->mutex);
	crossbar->simulation_done = 1;
	if (crossbar->result_count < crossbar->config.num_timesteps) {
		fail_locked(crossbar);
	}
	pthread_cond_broadcast(&crossbar->progress);
	pthread_cond_broadcast(&crossbar->state_available);
	pthread_mutex_unlock(&crossbar->mutex);
	return 0;
}

// gets the crossbar row number from an ngspice source name
static int row_from_name(const char *name, size_t *row_out)
{
	const char *digits = name + strlen(name);
	while (digits > name && isdigit((unsigned char)digits[-1])) {
		digits--;
	}
	if (*digits == '\0') {
		return -1;
	}
	char *end = NULL;
	unsigned long value = strtoul(digits, &end, 10);
	if (end == digits || *end != '\0') {
		return -1;
	}
	*row_out = (size_t)value;
	return 0;
}

// converts the current ngspice simulation time into a reservoir timestep
static size_t timestep_for_time(const Online_Crossbar *crossbar, double time)
{
	if (time <= 0.0) {
		return 0;
	}
	double scaled = time / crossbar->config.time_step;
	/* Hold state[t] over the full interval (t*dt, (t+1)*dt]. */
	double interval = ceil(scaled - 1.0e-9);
	size_t timestep = interval <= 1.0 ? 0 : (size_t)interval - 1;
	if (timestep >= crossbar->config.num_timesteps) {
		timestep = crossbar->config.num_timesteps - 1;
	}
	return timestep;
}

// gives ngspice the reservoir voltage for a requested crossbar row
static int callback_voltage(double *voltage, double time, char *name, int ident,
			    void *user_data)
{
	(void)ident;
	Online_Crossbar *crossbar = user_data;
	size_t row;
	if (!name || row_from_name(name, &row) != 0 ||
	    row >= crossbar->config.num_neurons) {
		fprintf(stderr, "ngspice: unrecognized external source %s\n",
			name ? name : "(null)");
		return 1;
	}

	size_t timestep = timestep_for_time(crossbar, time);
	pthread_mutex_lock(&crossbar->mutex);
	while (!crossbar->failed && crossbar->published_count <= timestep) {
		pthread_cond_wait(&crossbar->state_available, &crossbar->mutex);
	}
	if (crossbar->failed) {
		pthread_mutex_unlock(&crossbar->mutex);
		return 1;
	}

	*voltage =
	    crossbar->config.spike_amplitude *
	    crossbar->state_ring[(timestep % 2) * crossbar->config.num_neurons +
				 row];
	if (crossbar->consumed_count < timestep + 1) {
		crossbar->consumed_count = timestep + 1;
		pthread_cond_broadcast(&crossbar->progress);
	}
	pthread_mutex_unlock(&crossbar->mutex);
	return 0;
}

// collects and decodes crossbar output voltages at each timestep
static int callback_data(pvecvaluesall values, int count, int ident,
			 void *user_data)
{
	(void)count;
	(void)ident;
	Online_Crossbar *crossbar = user_data;
	double time = -1.0;
	double *columns =
	    calloc(crossbar->config.num_outputs * 2, sizeof(*columns));
	unsigned char *found =
	    calloc(crossbar->config.num_outputs * 2, sizeof(*found));
	if (!columns || !found) {
		free(columns);
		free(found);
		pthread_mutex_lock(&crossbar->mutex);
		fail_locked(crossbar);
		pthread_mutex_unlock(&crossbar->mutex);
		return 1;
	}

	for (int i = 0; i < values->veccount; i++) {
		pvecvalues value = values->vecsa[i];
		if (value->is_scale || strcmp(value->name, "time") == 0) {
			time = value->creal;
			continue;
		}
		const char *col = strstr(value->name, "col");
		if (!col) {
			continue;
		}
		char *end = NULL;
		unsigned long index = strtoul(col + 3, &end, 10);
		if (end != col + 3 &&
		    index < crossbar->config.num_outputs * 2) {
			columns[index] = value->creal;
			found[index] = 1;
		}
	}

	if (time >= 0.0) {
		double scaled = time / crossbar->config.time_step;
		double rounded = nearbyint(scaled);
		if (fabs(scaled - rounded) <= 1.0e-7 && rounded >= 1.0 &&
		    (size_t)rounded <= crossbar->config.num_timesteps) {
			size_t timestep = (size_t)rounded - 1;
			int complete = 1;
			for (size_t i = 0; i < crossbar->config.num_outputs * 2;
			     i++) {
				complete = complete && found[i];
			}
			if (complete) {
				double *decoded =
				    crossbar->results +
				    timestep * crossbar->config.num_outputs;
				if (convert_output_to_software(
					crossbar->config.num_neurons,
					crossbar->config.num_outputs, 1,
					columns, crossbar->resistances,
					crossbar->config.load_resistance,
					&crossbar->mapping, NULL,
					crossbar->config.spike_amplitude,
					decoded) == 0) {
					pthread_mutex_lock(&crossbar->mutex);
					if (crossbar->result_count <
					    timestep + 1) {
						crossbar->result_count =
						    timestep + 1;
					}
					if (timestep + 1 <
					    crossbar->config.num_timesteps) {
						ngSpice_SetBkpt(
						    (timestep + 2) *
						    crossbar->config.time_step);
					}
					pthread_cond_broadcast(
					    &crossbar->progress);
					pthread_mutex_unlock(&crossbar->mutex);
				}
			}
		}
	}

	free(columns);
	free(found);
	return 0;
}

// allocates the online crossbar and generates its ngspice netlist
int online_crossbar_init(const Online_Crossbar_Config *config,
			 const spires_reservoir *reservoir,
			 Online_Crossbar **crossbar_out)
{
	if (!crossbar_out) {
		return -1;
	}
	*crossbar_out = NULL;
	if (!reservoir || !config_is_valid(config)) {
		return -1;
	}

	Online_Crossbar *crossbar = calloc(1, sizeof(*crossbar));
	if (!crossbar) {
		return -1;
	}
	crossbar->config = *config;
	if (pthread_mutex_init(&crossbar->mutex, NULL) != 0) {
		goto fail;
	}
	crossbar->mutex_initialized = 1;
	if (pthread_cond_init(&crossbar->state_available, NULL) != 0) {
		goto fail;
	}
	crossbar->state_condition_initialized = 1;
	if (pthread_cond_init(&crossbar->progress, NULL) != 0) {
		goto fail;
	}
	crossbar->progress_condition_initialized = 1;

	crossbar->state_ring = calloc(2 * config->num_neurons, sizeof(double));
	crossbar->results =
	    calloc(config->num_timesteps * config->num_outputs, sizeof(double));
	if (!crossbar->state_ring || !crossbar->results) {
		goto fail;
	}
	if (convert_weights_to_resistances(
		reservoir, config->num_neurons, config->num_outputs,
		config->r_on, config->r_off, &crossbar->resistances,
		&crossbar->mapping) != 0) {
		goto fail;
	}

	Crossbar_Config netlist = {.rows = config->num_neurons,
				   .columns = config->num_outputs * 2,
				   .input_series = crossbar->state_ring,
				   .num_samples = config->num_timesteps,
				   .initial_resistance = crossbar->resistances,
				   .model_path = config->model_path,
				   .subcircuit_name = config->subcircuit_name,
				   .load_resistance = config->load_resistance,
				   .time_step = config->time_step,
				   .stop_time = config->num_timesteps *
						config->time_step,
				   .print_state_nodes = 0};
	if (generate_online_crossbar(config->netlist_path, &netlist) != 0) {
		goto fail;
	}

	*crossbar_out = crossbar;
	return 0;

fail:
	online_crossbar_destroy(crossbar);
	return -1;
}

// starts the persistent ngspice simulation in the background
int online_crossbar_start(Online_Crossbar *crossbar)
{
	if (!crossbar || crossbar->started) {
		return -1;
	}
	if (ngSpice_Init(callback_text, callback_text, callback_exit,
			 callback_data, callback_init_data, callback_background,
			 crossbar) != 0) {
		return -1;
	}
	if (ngSpice_Init_Sync(callback_voltage, NULL, NULL, NULL, crossbar) !=
	    0) {
		ngSpice_Reset();
		return -1;
	}

	char command[4096];
	int written = snprintf(command, sizeof(command), "source %s",
			       crossbar->config.netlist_path);
	if (written < 0 || (size_t)written >= sizeof(command) ||
	    ngSpice_Command(command) != 0) {
		ngSpice_Reset();
		return -1;
	}
	crossbar->started = 1;
	ngSpice_SetBkpt(crossbar->config.time_step);
	if (ngSpice_Command("bg_run") != 0) {
		ngSpice_Reset();
		crossbar->started = 0;
		return -1;
	}
	return 0;
}

// submits one reservoir state and returns the previous timestep's output
int online_crossbar_submit(Online_Crossbar *crossbar, size_t timestep,
			   const double *state, double *previous_output,
			   int *output_ready)
{
	if (!crossbar || !state || !output_ready || !crossbar->started ||
	    timestep >= crossbar->config.num_timesteps ||
	    timestep != crossbar->published_count ||
	    (timestep != 0 && !previous_output)) {
		return -1;
	}

	pthread_mutex_lock(&crossbar->mutex);
	while (!crossbar->failed && timestep >= crossbar->consumed_count + 2) {
		pthread_cond_wait(&crossbar->progress, &crossbar->mutex);
	}
	if (crossbar->failed) {
		pthread_mutex_unlock(&crossbar->mutex);
		return -1;
	}
	memcpy(crossbar->state_ring +
		   (timestep % 2) * crossbar->config.num_neurons,
	       state, crossbar->config.num_neurons * sizeof(double));
	crossbar->published_count++;
	pthread_cond_broadcast(&crossbar->state_available);

	*output_ready = timestep != 0;
	if (timestep != 0) {
		while (!crossbar->failed && crossbar->result_count < timestep) {
			pthread_cond_wait(&crossbar->progress,
					  &crossbar->mutex);
		}
		if (crossbar->failed) {
			pthread_mutex_unlock(&crossbar->mutex);
			return -1;
		}
		memcpy(previous_output,
		       crossbar->results +
			   (timestep - 1) * crossbar->config.num_outputs,
		       crossbar->config.num_outputs * sizeof(double));
	}
	pthread_mutex_unlock(&crossbar->mutex);
	return 0;
}

// waits for and returns the final crossbar output
int online_crossbar_finish(Online_Crossbar *crossbar, double *final_output)
{
	if (!crossbar || !final_output ||
	    crossbar->published_count != crossbar->config.num_timesteps) {
		return -1;
	}
	pthread_mutex_lock(&crossbar->mutex);
	while (!crossbar->failed &&
	       crossbar->result_count < crossbar->config.num_timesteps) {
		pthread_cond_wait(&crossbar->progress, &crossbar->mutex);
	}
	if (crossbar->failed) {
		pthread_mutex_unlock(&crossbar->mutex);
		return -1;
	}
	memcpy(final_output,
	       crossbar->results + (crossbar->config.num_timesteps - 1) *
				       crossbar->config.num_outputs,
	       crossbar->config.num_outputs * sizeof(double));
	pthread_mutex_unlock(&crossbar->mutex);
	return 0;
}

// stops ngspice and frees all online crossbar resources
void online_crossbar_destroy(Online_Crossbar *crossbar)
{
	if (!crossbar) {
		return;
	}
	if (crossbar->started && ngSpice_running()) {
		ngSpice_Command("bg_halt");
		if (crossbar->mutex_initialized) {
			pthread_mutex_lock(&crossbar->mutex);
			while (!crossbar->simulation_done &&
			       ngSpice_running()) {
				pthread_cond_wait(&crossbar->progress,
						  &crossbar->mutex);
			}
			pthread_mutex_unlock(&crossbar->mutex);
		}
	}
	if (crossbar->started) {
		ngSpice_Reset();
	}
	free(crossbar->resistances);
	free(crossbar->state_ring);
	free(crossbar->results);
	if (crossbar->state_condition_initialized) {
		pthread_cond_destroy(&crossbar->state_available);
	}
	if (crossbar->progress_condition_initialized) {
		pthread_cond_destroy(&crossbar->progress);
	}
	if (crossbar->mutex_initialized) {
		pthread_mutex_destroy(&crossbar->mutex);
	}
	free(crossbar);
}