#include "diffeq.h"

#include <string.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>

/* Calling:
 * if data is non-null, t is used. if data is null, t is not used.
 *
 * 1. Initializing files: files = [array of files], data = NULL, n = sizeof [array of files]
 * 2. Updating files: files = anything non-NULL [unused], data = [array of data], n = sizeof [array of data]. 
 * 	n may be 0, then saved n from init is used
 * 3. Closing files: files = NULL, data = NULL, n = anything 
 *
 * returns: whether the operation was successful
 */
char flush_approx(const char ** files, f64 t, f64 * data, usz n) {
	static FILE ** handlers;
	static usz size;
	usz i;

	if (data == NULL && files == NULL) {
		for (i = 0; i < n; ++i) {
			if (fclose(handlers[i]) != 0)
				perror("closing files");
		}
		free(handlers);
		return 1;
	}
	if (data == NULL) {
		size = n;
		handlers = calloc(n, sizeof (FILE*));
		for (i = 0; i < n; ++i) {
			handlers[i] = fopen(files[i], "w+");
			if (handlers[i] == NULL) {
				perror(files[i]);
				return 0;
			}
		}
		return 1;
	}
	if (files == NULL) {
		for (i = 0; i < (n == 0 ? size : n); ++i)
			fprintf(handlers[i], "%f %f\n", t, data[i]);
		return 1;
	}
	return 0;
}

void difeq_solve_euler(eqf_t *eqfs, f64 *init, const char **files, usz n, f64 a, f64 b, f64 dt) {
	f64 *guesses1 = calloc(n, sizeof(f64));
	f64 *guesses2 = calloc(n, sizeof(f64));
	f64 cur = a;
	size_t i;

	flush_approx(files, 0, NULL, n);

	memcpy(guesses1, init, n * sizeof(f64));
	memcpy(guesses2, init, n * sizeof(f64));
	
	for(cur = a; cur < b; cur += dt) {
		for (i = 0; i < n; ++i)
			guesses2[i] += dt * eqfs[i](guesses1, n, cur); 
		memcpy(guesses1, guesses2, n * sizeof(f64)); 
		flush_approx(NULL, cur, guesses1, n);
	}
	flush_approx(NULL, 0, NULL, 0);
	free(guesses1);
	free(guesses2);
}

void difeq_solve_RK(eqf_t *eqfs, f64 *init, const char **files, usz n, f64 a, f64 b, f64 dt) {
	f64 *guesses1 = calloc(n, sizeof(f64));
	f64 *guesses2 = calloc(n, sizeof(f64));
	f64 cur = a;
	f64 *k1s = calloc(n, sizeof(f64)), 
	    *k2s = calloc(n, sizeof(f64)), 
	    *k3s = calloc(n, sizeof(f64)), 
	    *k4s = calloc(n, sizeof(f64));
	size_t step, i;

	flush_approx(files, 0, NULL, n);

	memcpy(guesses2, init, n * sizeof(f64));
	memcpy(guesses1, init, n * sizeof(f64));
	
	for(step = 0; step < (b-a)/dt; cur = a + ++step*dt) {
		for (i = 0; i < n; ++i)
			k1s[i] = eqfs[i](guesses1, n, cur);
		for (i = 0; i < n; ++i)
			guesses2[i] = guesses1[i] + dt/2 * k1s[i];
		for (i = 0; i < n; ++i)
			k2s[i] = eqfs[i](guesses2, n, cur + dt/2);
		for (i = 0; i < n; ++i)
			guesses2[i] = guesses1[i] + dt/2 * k2s[i];
		for (i = 0; i < n; ++i)
			k3s[i] = eqfs[i](guesses2, n, cur + dt/2);
		for (i = 0; i < n; ++i)
			guesses2[i] = guesses1[i] + dt * k3s[i];
		for (i = 0; i < n; ++i)
			k4s[i] = eqfs[i](guesses2, n, cur + dt);

		for (i = 0; i < n; ++i)
			guesses2[i] += dt/6 * (k1s[i] + 2*k2s[i] +2*k3s[i] + k4s[i]); 

		memcpy(guesses1, guesses2, n * sizeof(f64)); 
		flush_approx(NULL, cur, guesses1, n);
	}
	flush_approx(NULL, 0, NULL, 0);
	free(guesses1);
	free(guesses2);
	free(k1s);
	free(k2s);
	free(k3s);
	free(k4s);
}