#include "spires_interface.h"
#include <spires.h>

#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

int collect_reservoir_states(spires_reservoir *reservoir,
			     const double *input_series, size_t series_length,
			     Reservoir_State_Matrix *result)
{
	if (!reservoir || !input_series || series_length == 0 || !result) {
		return -1;
	}

	result->num_samples = 0;
	result->num_features = 0;
	result->states = NULL;

	const size_t num_inputs = spires_num_inputs(reservoir);
	const size_t num_neurons = spires_num_neurons(reservoir);
	if (num_inputs == 0 || num_neurons == 0) {
		return -1;
	}

	if (series_length > SIZE_MAX / num_neurons ||
	    series_length * num_neurons > SIZE_MAX / sizeof(double)) {
		fprintf(stderr, "Reservoir state matrix is too large\n");
		return -1;
	}

	double *states = malloc(num_neurons * series_length * sizeof(*states));
	if (!states) {
		fprintf(stderr, "Failed to allocate reservoir state matrix\n");
		return -1;
	}

	spires_status status = spires_reservoir_reset(reservoir);
	if (status != SPIRES_OK) {
		fprintf(stderr, "Failed to reset SPIRES reservoir\n");
		free(states);
		return -1;
	}

	for (size_t timestep = 0; timestep < series_length; timestep++) {
		const double *current_input =
		    input_series + timestep * num_inputs;
		status = spires_step(reservoir, current_input);
		if (status != SPIRES_OK) {
			fprintf(stderr, "SPIRES step failed at timestep %zu\n",
				timestep);
			free(states);
			return -1;
		}

		double *current_state = states + timestep * num_neurons;
		status = spires_read_reservoir_state(reservoir, current_state);
		if (status != SPIRES_OK) {
			fprintf(
			    stderr,
			    "Failed to read reservoir state at timestep %zu\n",
			    timestep);
			free(states);
			return -1;
		}
	}

	result->num_samples = series_length;
	result->num_features = num_neurons;
	result->states = states;

	return 0;
}

int convert_weights_to_resistances(const spires_reservoir *reservoir,
				   size_t num_neurons, size_t num_outputs,
				   double r_on, double r_off,
				   double **resistances_out,
				   conductance_mapping *mapping)
{
	if (!reservoir || num_neurons == 0 || num_outputs == 0 || r_on <= 0.0 ||
	    r_off <= r_on || !resistances_out || !mapping) {
		return -1;
	}

	*resistances_out = NULL;
	if (num_outputs > SIZE_MAX / 2 ||
	    num_neurons > SIZE_MAX / num_outputs) {
		return -1;
	}

	size_t weight_count = num_neurons * num_outputs;
	size_t num_physical_columns = num_outputs * 2;
	if (num_neurons > SIZE_MAX / num_physical_columns ||
	    num_neurons * num_physical_columns > SIZE_MAX / sizeof(double)) {
		return -1;
	}

	double *readout = malloc(weight_count * sizeof(*readout));
	if (!readout) {
		fprintf(stderr, "Failed to allocate readout weights\n");
		return -1;
	}
	if (spires_read_readout(reservoir, readout) != SPIRES_OK) {
		fprintf(stderr, "Failed to read SPIRES readout weights\n");
		free(readout);
		return -1;
	}

	double *resistances =
	    malloc(num_neurons * num_physical_columns * sizeof(*resistances));
	if (!resistances) {
		fprintf(stderr, "Failed to allocate crossbar resistances\n");
		free(readout);
		return -1;
	}

	double max_abs_weight = 0.0;
	for (size_t i = 0; i < weight_count; i++) {
		if (!isfinite(readout[i])) {
			fprintf(stderr,
				"Invalid readout weight at index %zu: %g\n", i,
				readout[i]);
			free(readout);
			free(resistances);
			return -1;
		}

		double abs_weight = fabs(readout[i]);
		if (abs_weight > max_abs_weight) {
			max_abs_weight = abs_weight;
		}
	}
	if (max_abs_weight == 0.0) {
		fprintf(stderr,
			"Cannot map an all-zero readout to resistances\n");
		free(readout);
		free(resistances);
		return -1;
	}

	mapping->g_min = 1.0 / r_off;
	mapping->g_max = 1.0 / r_on;
	mapping->max_abs_weight = max_abs_weight;

	mapping->alpha = (mapping->g_max - mapping->g_min) / max_abs_weight;

	for (size_t neuron = 0; neuron < num_neurons; neuron++) {
		for (size_t output = 0; output < num_outputs; output++) {
			size_t weight_index = output * num_neurons + neuron;
			size_t positive_column = 2 * output;
			size_t negative_column = positive_column + 1;
			double positive_conductance;
			double negative_conductance;

			size_t positive_index =
			    neuron * num_physical_columns + positive_column;
			size_t negative_index =
			    neuron * num_physical_columns + negative_column;

			double weight = readout[weight_index];
			if (weight >= 0.0) {
				positive_conductance =
				    mapping->g_min + mapping->alpha * weight;
				negative_conductance = mapping->g_min;
			} else {
				positive_conductance = mapping->g_min;
				negative_conductance =
				    mapping->g_min + mapping->alpha * (-weight);
			}

			resistances[positive_index] =
			    1.0 / positive_conductance;
			resistances[negative_index] =
			    1.0 / negative_conductance;
		}
	}

	*resistances_out = resistances;
	free(readout);

	return 0;
}

int train_reservoir(spires_reservoir *reservoir, double *input_series,
		    double *target_series, size_t series_length, double lambda)
{
	if (!reservoir || !input_series || !target_series ||
	    series_length == 0 || lambda < 0.0) {
		return -1;
	}

	spires_status status = spires_train_ridge(
	    reservoir, input_series, target_series, series_length, lambda);

	if (status != SPIRES_OK) {
		fprintf(stderr, "SPIRES ridge training failed\n");
		return -1;
	}

	return 0;
}

void free_reservoir_state_matrix(Reservoir_State_Matrix *matrix)
{
	if (!matrix) {
		return;
	}

	free(matrix->states);
	matrix->states = NULL;
	matrix->num_samples = 0;
	matrix->num_features = 0;
}