#include "benchmark.h"

#include <dirent.h>
#include <math.h>
#include <spires.h>
#include <stdio.h>
#include <stdlib.h>

// spires reservoir parameters
#define NUM_NEURONS 600
#define NUM_INPUTS 1
#define NUM_OUTPUTS 1
#define SPECTRAL_RADIUS 0.95
#define EI_RATIO 0.8
#define INPUT_STRENGTH 0.1
#define CONNECTIVITY 0.1
#define DT 1.0

// ridge regression training
#define PI 3.14159265358979323846
#define LAMBDA 1.0e-4

#define NUM_TRAINING_STEPS 500
#define NUM_CROSSBAR_COLUMNS (NUM_OUTPUTS * 2)
// #define NUM_STEPS 2000

int main(void)
{
	/* ---------- LIST ALL MODELS HERE ----------*/
	MemModel models[] = {
	    {
		.model_path = "models/fixed_resistor.cir",
		.subcircuit_name = "fixed_resistor",
	    },
	    {
		.model_path = "models/hp_memristor.cir",
		.subcircuit_name = "memristor",
	    },
	    {
		.model_path = "models/yakopcic_memristor.cir",
		.subcircuit_name = "MEM_YAKOPCIC",
	    },
	    {
		.model_path = "models/Pershin_DiVentra_memristor.cir",
		.subcircuit_name = "memristor",
	    },
	};

	size_t model_count = sizeof(models) / sizeof(models[0]);

	/* ---------- SPIRES SET UP ----------*/
	// discrete LIF parameters for spires
	double lif_config[] = {
	    0.0, // V_off
	    1.0, // V_th
	    0.2, // leak rate
	    0.5, // bias
	};

	const spires_reservoir_config config = {
	    .num_neurons = NUM_NEURONS,
	    .num_inputs = NUM_INPUTS,
	    .num_outputs = NUM_OUTPUTS,
	    .spectral_radius = SPECTRAL_RADIUS,
	    .ei_ratio = EI_RATIO,
	    .input_strength = INPUT_STRENGTH,
	    .connectivity = CONNECTIVITY,
	    .dt = DT,
	    .connectivity_type = SPIRES_CONN_RANDOM,
	    .neuron_type = SPIRES_NEURON_LIF_DISCRETE,
	    .neuron_params = lif_config};

	spires_reservoir *reservoir = NULL;
	if (spires_reservoir_create(&config, &reservoir) != 0) {
		fprintf(stderr, "Failed to create reservoir");
		return -1;
	}

	/* ---------- Train spires readout layer ----------*/
	/* ---------- Training inputs ----------*/
	double training_inputs[NUM_TRAINING_STEPS * NUM_INPUTS];
	for (size_t timestep = 0; timestep < NUM_TRAINING_STEPS; timestep++) {
		for (size_t input = 0; input < NUM_INPUTS; input++) {
			double signal =
			    0.7 * sin(2.0 * PI * (double)timestep / 50.0) +
			    0.3 * sin(2.0 * PI * (double)timestep / 17.0);
			training_inputs[timestep * NUM_INPUTS + input] = signal;
		}
	}

	/* ---------- Target outputs  ----------*/
	double target_outputs[NUM_TRAINING_STEPS * NUM_OUTPUTS];
	for (size_t timestep = 0; timestep < NUM_TRAINING_STEPS; timestep++) {
		size_t next_timestep = (timestep + 1) % NUM_TRAINING_STEPS;

		double target =
		    0.7 * sin(2.0 * PI * (double)next_timestep / 50.0) +
		    0.3 * sin(2.0 * PI * (double)next_timestep / 17.0);
		for (size_t output = 0; output < NUM_OUTPUTS; output++) {
			target_outputs[timestep * NUM_OUTPUTS + output] =
			    target;
		}
	}

	/* ---------- Ridge Regression ----------*/
	if (train_reservoir(reservoir, training_inputs, target_outputs,
			    NUM_TRAINING_STEPS, LAMBDA) < 0) {
		fprintf(stderr, "Failed to train the spires reservoir");
		spires_reservoir_destroy(reservoir);
		return -1;
	}

	/* ---------- Collect reservoir states ----------*/
	Reservoir_State_Matrix state_matrix = {0};
	if (collect_reservoir_states(reservoir, training_inputs,
				     NUM_TRAINING_STEPS, &state_matrix) != 0) {
		fprintf(stderr, "Failed to collect reservoir states");
		spires_reservoir_destroy(reservoir);
		return -1;
	}
	printf("collected state matrix size: %zu x %zu\n",
	       state_matrix.num_samples, state_matrix.num_features);

	/* ---------- Generate raster plot ----------*/
	if (plot_raster(&state_matrix, NUM_NEURONS, 0.5) != 0) {
		fprintf(stderr, "Failed to plot raster");
	}

	/* ---------- Run Benchmark on each model ----------*/
	size_t predictions_per_model =
	    state_matrix.num_samples * config.num_outputs;

	double *predictions = malloc(model_count * NUM_OUTPUTS *
				     NUM_TRAINING_STEPS * sizeof(double));
	if (predictions == NULL) {
		fprintf(stderr, "Failed to allocate memroy for predictions");
		free_reservoir_state_matrix(&state_matrix);
		spires_reservoir_destroy(reservoir);
		return -1;
	}

	double mean_squared_error[model_count];

	for (size_t model = 0; model < model_count; model++) {
		double *model_predictions =
		    predictions + model * predictions_per_model;
		printf("\n\nRunning benchmark on %s\n",
		       models[model].model_path);

		if (run_benchmark(&config, reservoir, &state_matrix,
				  models[model].model_path,
				  models[model].subcircuit_name,
				  model_predictions) < 0) {
			fprintf(stderr, "Failed to run benchmark");
			free(predictions);
			free_reservoir_state_matrix(&state_matrix);
			spires_reservoir_destroy(reservoir);
			return -1;
		}

		plot_reservoir_predictions(
		    target_outputs, model_predictions, state_matrix.num_samples,
		    config.num_outputs, 0, models[model].model_path);

		mean_squared_error[model] =
		    calculate_MSE(target_outputs, model_predictions,
				  state_matrix.num_samples, config.num_outputs);
	}

	for (size_t model = 0; model < model_count; model++) {
		printf("Model: %s, MSE: %.17g\n", models[model].model_path,
		       mean_squared_error[model]);
	}

	free(predictions);
	free_reservoir_state_matrix(&state_matrix);
	spires_reservoir_destroy(reservoir);

	return 0;
}