#include "common.h"

#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <math.h>

#ifdef _GNU_SOURCE
#undef _GNU_SOURCE
#endif

struct config {
	const char         *filename;
	bool                play_audio;     /* If true, this program will play the samples */
	double              gain;
	unsigned long       offset_ms;      /* Time offset in ms for plotting */
};

struct audio_data {
	size_t              nbytes;
	size_t              n_sam;          /* Number of samples */
	unsigned long       sample_rate;    /* Sample rate in Hz */
	double             *orig_audio;     /* Original audio sample */
	double             *ampl_audio;     /* Amplified/clipped audio sample */
	double             *fixd_audio;     /* Fixed/restored audio sample */
};

const char *octave_fft_payload =
	"ds_attY = fft(signal);\n" /* Fast fourier transform */
	"ss_attY_half = ds_attY(1:n_sam/2+1);\n"
	"ss_attY_half(2:end-1) = 2 * ss_attY_half(2:end-1);\n"
	"df = Fs/n_sam;\n"
	"vf_lo = floor(80/df);\n"
	"vf_hi = floor(500/df);\n"
	"fc_btm = floor(fc/df);\n"
	"fc_top = size(ss_attY_half, 1);\n"
	"ss_fixed_half = vertcat(\n"
	"  ss_attY_half(1:vf_lo) * 0,\n" /* Lower cutoff */
	"  ss_attY_half(vf_lo+1:vf_hi),\n"
	"  ss_attY_half(vf_hi+1:fc_btm),\n"
	"  ss_attY_half(fc_btm+1:fc_top) * 0\n" /* Upper cutoff */
	");\n"
	"ds_fixed = vertcat(\n"
	"  ss_fixed_half(1),\n"
	"  0.5 * ss_fixed_half(2:end-1),\n"
	"  ss_fixed_half(end),\n"
	"  fliplr(0.5 * ss_fixed_half(2:end-1)),\n"
	"  ss_fixed_half(1)\n"
	");\n"
	"fixed_signal = real(ifft(ds_fixed));\n"; /* Real part of the ifft() */

const char *octave_plot_payload =
	"orig_audio = samples_subset(:, 1);\n" /* Extract column vectors */
	"ampl_audio = samples_subset(:, 2);\n"
	"fixd_audio = samples_subset(:, 3);\n"
	"figure\n"
	"subplot(3, 1, 1);\n" /* First plot */
	"plot(t, orig_audio);\n"
	"title('Input signal');\n"
	"ax1 = gca();\n"
	"set(ax1, 'xgrid', 'on');\n"
	"set(ax1, 'ygrid', 'on');\n"
	"set(ax1, 'xlabel', 't [ms]');\n"
	"set(ax1, 'ylabel', 'Amplitude [-]');\n"
	"subplot(3, 1, 2);\n"
	"plot(t, ampl_audio);\n" /* Second plot */
	"title('Amplified signal');\n"
	"ax2 = gca();\n"
	"set(ax2, 'xgrid', 'on');\n"
	"set(ax2, 'ygrid', 'on');\n"
	"set(ax2, 'xlabel', 't [ms]');\n"
	"set(ax2, 'ylabel', 'Amplitude [-]');\n"
	"subplot(3, 1, 3);\n"
	"plot(t, fixd_audio);\n"  /* Third plot */
	"title('Fixed audio signal');\n"
	"ax3 = gca();\n"
	"set(ax3, 'xgrid', 'on');\n"
	"set(ax3, 'ygrid', 'on');\n"
	"set(ax3, 'xlabel', 't [ms]');\n"
	"set(ax3, 'ylabel', 'Amplitude [-]');\n";

void free_samples(struct audio_data *aud);
int alloc_samples(struct audio_data *aud);
int load_audio_sample(struct config *cfg, struct audio_data *aud);
void display_graphs(struct config *cfg, struct audio_data *aud, FILE *fp_oct);
int play_samples(struct audio_data *aud);
double calc_error(struct audio_data *aud);
double rms(double *vect, size_t n);

static inline unsigned long offset_to_index(unsigned long offset_ms,
	unsigned long sample_rate);
static void print_usage(const char *app_name);
static void parse_args(int argc, const char **argv, struct config *cfg);

/* Root mean square */
double rms(double *vect, size_t n)
{
	size_t i;
	double sum = 0.0;

	for(i = 0; i < n; i++)
		sum += vect[i] * vect[i];

	return sqrt(sum / n);
}

/* Calculate the error between the original audio sample and the restored sample */
double calc_error(struct audio_data *aud)
{
	size_t i;
	double diff_sq = 0;
	double orig_sq = 0;

	if (!aud) {
		print_error("Critical error: null pointer dereference at:\n");
		print_error(" calc_error()\n");
		raise(SIGSEGV);
	}

	for (i = 0; i < aud->n_sam; i++) {
		double si_orig = aud->orig_audio[i] * aud->orig_audio[i];
		double si_fixd = aud->fixd_audio[i] * aud->fixd_audio[i];

		diff_sq += si_orig - si_fixd;
		orig_sq += si_orig;
	}

	if (orig_sq == 0)
		return 0;

	return diff_sq/orig_sq;
}

/* Plays all 3 audio samples in the `audio_data` struct using `aplay` */
int play_samples(struct audio_data *aud)
{
	size_t i;
	int16_t *tmp;

	FILE *fp;
	char tmpfile[128] = {0};
	char cmdbuf[256] = {0};

	if (!aud) {
		print_error("Critical error: null pointer dereference at:\n");
		print_error(" calc_error()\n");
		raise(SIGSEGV);
	}

	/* Create temporary filename based on this process' pid */
	snprintf(tmpfile, sizeof(tmpfile), "/tmp/snd_orig_%d.raw", getpid());
	snprintf(cmdbuf, sizeof(cmdbuf), "aplay -f S16_LE -r 8000 %s", tmpfile);

	/* Allocate a temporary buffer */
	tmp = malloc(aud->nbytes * sizeof *tmp);
	if (!tmp)
		return -ENOMEM;

	fp = fopen(tmpfile, "w");

	/* Original audio */
	printf(COLOR_BLUE "Playing original sample..." COLOR_MAGENTA "\n");
	for (i = 1; i < aud->n_sam; i++) {
		tmp[i] = round(aud->orig_audio[i]);
	}
	fwrite(tmp, sizeof(int16_t), aud->n_sam, fp);
	fflush(fp);
	system(cmdbuf);

	fseek(fp, 0L, SEEK_SET);
	/* Clipped audio */
	printf(COLOR_BLUE "Playing clipped sample..." COLOR_MAGENTA "\n");
	for (i = 1; i < aud->n_sam; i++) {
		tmp[i] = round(aud->ampl_audio[i]);
	}
	fwrite(tmp, sizeof(int16_t), aud->n_sam, fp);
	fflush(fp);
	system(cmdbuf);

	fseek(fp, 0L, SEEK_SET);
	/* Fixed audio */
	printf(COLOR_BLUE "Playing clipped sample..." COLOR_MAGENTA "\n");
	for (i = 1; i < aud->n_sam; i++) {
		tmp[i] = round(aud->fixd_audio[i]);
	}
	fwrite(tmp, sizeof(int16_t), aud->n_sam, fp);
	fflush(fp);
	system(cmdbuf);

	printf(COLOR_OFF "\n");

	/* Cleanup and remove the temporary file */
	fclose(fp);
	remove(tmpfile);
	free(tmp);
	return 0;
}

static inline unsigned long offset_to_index(unsigned long offset_ms, unsigned long sample_rate)
{
	return (unsigned long)(offset_ms * (double)sample_rate/1000.0);
}

/*
 * This function creates 3 plots of amplitude vs time, for every audio data in
 * the `audio_data` struct, each one of 40 [ms] long starting from the offset
 * given in the `config` struct.
 */
void display_graphs(struct config *cfg, struct audio_data *aud, FILE *fp_oct)
{
	size_t i;
	unsigned long lower, upper;

	if (!cfg || !aud) {
		print_error("Critical error: null pointer dereference at:\n");
		print_error(" display_graphs()\n");
		raise(SIGSEGV);
	}

	lower = offset_to_index(cfg->offset_ms, aud->sample_rate);
	upper = offset_to_index(cfg->offset_ms + 40, aud->sample_rate);

	if (lower >= aud->n_sam || upper >= aud->n_sam) {
		print_warn("Warning: offset is out of boundary, fixing to view the last 40 [ms]\n");
		upper = aud->n_sam - 1;
		lower = upper - offset_to_index(40, aud->sample_rate);
	}

	fprintf(fp_oct, "clear all\n");
	fprintf(fp_oct, "Fs = %lu;\n", aud->sample_rate);
	fprintf(fp_oct, "dt = 1000/Fs;\n");
	fprintf(fp_oct, "t = dt * (%lu:1:%lu);\n", lower, upper);
	fflush(fp_oct);

	fprintf(fp_oct, "samples_subset = [");
	for (i = lower; i <= upper; i++) {
		fprintf(fp_oct, "%f %f %f; ",
			aud->orig_audio[i], aud->ampl_audio[i], aud->fixd_audio[i]);
	}
	fprintf(fp_oct, "];\n");
	fflush(fp_oct);

	fprintf(fp_oct, "%s", octave_plot_payload);
	fflush(fp_oct);
}

/*
 * This function attempts to restore a clipped audio sample
 * by sending specific instructions to GNU Octave (using pipes).
 */
int restore_audio(struct config *cfg, struct audio_data *aud, FILE *fp_oct)
{
	size_t i;
	char err_s[BUFSIZ];

	FILE *fp_tmp;
	char tmpfile[128] = {0};
	char line_buffer[128];

	if (!cfg || !aud) {
		print_error("Critical error: null pointer dereference at:\n");
		print_error(" restore_audio()\n");
		raise(SIGSEGV);
	}

	snprintf(tmpfile, sizeof(tmpfile), "tmp_%d", getpid());

	fprintf(fp_oct, "clear all\n");
	fflush(fp_oct);

	double rms_ratio;
	rms_ratio = rms(aud->orig_audio, aud->n_sam) /
	            rms(aud->ampl_audio, aud->n_sam);

	fprintf(fp_oct, "signal = [");
	for (i = 0; i < aud->n_sam; i++) {
		aud->fixd_audio[i] = aud->ampl_audio[i] * rms_ratio;
	//	aud->fixd_audio[i] = aud->ampl_audio[i] / cfg->gain;
		fprintf(fp_oct, "%f; ", aud->fixd_audio[i]);
	}
	fprintf(fp_oct, "];\n");
	fflush(fp_oct);

	fprintf(fp_oct, "gain = %f;\n", cfg->gain);
	fprintf(fp_oct, "Fs = %lu;\n", aud->sample_rate);
	fprintf(fp_oct, "fc = 3200;\n");
	fprintf(fp_oct, "n_sam = size(signal, 1);\n");
	fflush(fp_oct);

	fprintf(fp_oct, "%s", octave_fft_payload);
	fflush(fp_oct);

	fprintf(fp_oct, "csvwrite('%s', fixed_signal);\n", tmpfile);
	fflush(fp_oct);

	sleep(1);
	if (!(fp_tmp = fopen(tmpfile, "r"))) {
		strerror_r(errno, err_s, BUFSIZ);
		print_error("Error retrieving samples from GNU Octave: %s\n", err_s);
		return -errno;
	}

	for (i = 0; i < aud->n_sam; i++) {
		char *end;
		double sample = 0;

		if (fgets(line_buffer, sizeof(line_buffer), fp_tmp)) {
			sample = strtod(line_buffer, &end);
			if ((*end != '\0' && *end != '\n') || errno == ERANGE) {
				print_error("Warning: can't parse value from Octave <%s>\n", line_buffer);
				continue;
			}

			if (isgreaterequal(aud->ampl_audio[i], 32767) ||
				islessequal(aud->ampl_audio[i], -32768)) {
				aud->fixd_audio[i] = sample;
			}
		}
	}

	fclose(fp_tmp);
	remove(tmpfile);

	return 0;
}

/* Load input audio and perform the amplification/saturation */
int load_audio_sample(struct config *cfg, struct audio_data *aud)
{
	int ret;
	size_t i;
	struct stat buf;
	char err_s[BUFSIZ];

	int16_t *tmp;
	FILE *fp;

	if (!cfg || !aud) {
		print_error("Critical error: null pointer dereference at:\n");
		print_error(" load_audio_sample()\n");
		raise(SIGSEGV);
	}

	/* Check if we have read access to the input file and obtain its size */
	ret = stat(cfg->filename, &buf);
	if (ret < 0) {
		strerror_r(errno, err_s, BUFSIZ);
		print_error("Error opening file: %s\n", err_s);
		return -EACCES;
	}

	if (buf.st_size < 2) {
		print_error("Error: input file is empty (or too small)\n");
	}

	aud->nbytes = (buf.st_size % 2) ? (buf.st_size - 1) : buf.st_size;
	aud->n_sam = aud->nbytes / 2; /* fixed value (for now) */
	aud->sample_rate = 8000;

	/* Allocate a temporary buffer to read the file */
	tmp = malloc(aud->n_sam * sizeof *tmp);
	if (!tmp)
		return -ENOMEM;

	if (!(fp = fopen(cfg->filename, "rb"))) {
		strerror_r(errno, err_s, BUFSIZ);
		print_error("Error loading samples: %s\n", err_s);
		return -errno;
	}

	/*
	 * Load the entire file into the buffer, it reads
	 * `aud->n_sam` items of 2 bytes (16 bits) long.
	 */
	fread(tmp, sizeof(int16_t), aud->n_sam, fp);

	ret = alloc_samples(aud);
	if (ret)
		return ret;

	for (i = 0; i < aud->n_sam; i++) {
		double sam;
		aud->orig_audio[i] = tmp[i];

		/*
		 * Here we perform the saturation if the product goes
		 * out of `int16` range.
		 */
		sam = cfg->gain * tmp[i];

		if (sam < -32768)
			aud->ampl_audio[i] = -32768;
		else if (sam > 32767)
			aud->ampl_audio[i] = 32767;
		else
			aud->ampl_audio[i] = sam;
	}

	/* Cleanup */
	fclose(fp);
	free(tmp);

	return 0;
}

/* Frees the memory space */
void free_samples(struct audio_data *aud)
{
	if (!aud) {
		print_error("Critical error: null pointer dereference at\n");
		print_error(" alloc_samples()\n");
		raise(SIGSEGV);
	}

	if (aud->orig_audio)
		free(aud->orig_audio);
	if (aud->ampl_audio)
		free(aud->ampl_audio);
	if (aud->fixd_audio)
		free(aud->fixd_audio);

	/* Pointers are set to NULL to indicate deallocation */
	aud->orig_audio = aud->ampl_audio = aud->fixd_audio = NULL;
}

/* Allocate enough memory space to store samples */
int alloc_samples(struct audio_data *aud)
{
	if (!aud) {
		print_error("Critical error: null pointer dereference at\n");
		print_error(" alloc_samples()\n");
		raise(SIGSEGV);
	}

	aud->orig_audio = calloc(aud->n_sam, sizeof *aud->orig_audio);
	if (aud->orig_audio == NULL)
		goto oa_err;

	aud->ampl_audio = calloc(aud->n_sam, sizeof *aud->ampl_audio);
	if (aud->ampl_audio == NULL)
		goto aa_err;

	aud->fixd_audio = calloc(aud->n_sam, sizeof *aud->fixd_audio);
	if (aud->fixd_audio == NULL)
		goto fa_err;

	return 0;

fa_err:
	free(aud->ampl_audio);
aa_err:
	free(aud->orig_audio);
oa_err:
	print_error("Error: Out of memory\n");
	return -ENOMEM;
}

int main(int argc, const char **argv)
{
	int ret;
	unsigned int i;

	pid_t pid;
	int status;
	int pfd[2];
	FILE *fp;

	char err_s[BUFSIZ];

	struct audio_data _aud, *aud = &_aud;
	struct config _cfg, *cfg = &_cfg;
	parse_args(argc, argv, cfg);

	/* Load sample */
	if (load_audio_sample(cfg, aud) != 0)
		exit(-5);

	/* Create a pipe. */
	if (pipe(pfd) < 0) {
		strerror_r(errno, err_s, BUFSIZ);
		print_error("Error creating pipe: %s\n", err_s);
		exit(-3);
	}

	/* Create a child process. */
	debug_info("[%d] Creating child process...\n", getpid());
	if ((pid = fork()) < 0) {
		strerror_r(errno, err_s, BUFSIZ);
		print_error("Error spaning child process: %s\n", err_s);
		exit(-2);
	}

	/*
	 * Child process
	 */
	if (pid == 0) {
		int null_fd;
		int stdout_fd = dup(STDOUT_FILENO);
		int stderr_fd = dup(STDERR_FILENO);

		debug_info("[%d] Spawning GNU Octave...\n", getpid());

		/* Attach child's stdin to read from the pipe. */
		dup2(pfd[0], STDIN_FILENO);
		close(pfd[1]); /* Close the write end of the pipe */

		null_fd = open("/dev/null", O_WRONLY);
		dup2(null_fd, STDOUT_FILENO); /* Ignore any message to stdout */
		dup2(null_fd, STDERR_FILENO); /* Ignore any message to stderr */

		/* Replace the child's process image with GNU Octave */
		execlp("octave", "octave", "--interactive", NULL);

		/*
		 * The following lines will not be executed after a successful
		 * invocation of exec()
		 */
		strerror_r(errno, err_s, BUFSIZ);

		dup2(stdout_fd, STDOUT_FILENO); /* Restore stdout */
		dup2(stderr_fd, STDERR_FILENO); /* Restore stderr */

		print_error("Error executing GNU Octave: %s\n", err_s);
		exit(-1);
	}

	/* Close the pipe's writing end */
	close(pfd[0]);

	/* Wait for Octave to become idle */
	sleep(1);

	fp = fdopen(pfd[1], "w");
	restore_audio(cfg, aud, fp);

	if (cfg->play_audio)
		play_samples(aud);

	display_graphs(cfg, aud, fp);
	printf("Error: %.6f\n", calc_error(aud));

	printf("\nPress ENTER to exit...");
	getchar();

	fprintf(fp, "pause(1);\n");
	fprintf(fp, "exit;\n");
	fflush(fp);

	waitpid(pid, &status, 0);
	debug_info("[%d] GNU Octave returned: %d\n", getpid(), status);

	fclose(fp);
	free_samples(aud);

	return 0;
}

static void parse_args(int argc, const char **argv,
	struct config *cfg)
{
	char *end;
	cfg->play_audio = false;

	/* Checks the number of arguments and perform some input sanitization */
	switch (argc) {
	case 5:
		if (strcmp(argv[4], "p") == 0)
			cfg->play_audio = true;
		else
			print_warn("Warning: ignoring last argument...\n");
		/* Falls through */
	case 4:
		cfg->filename = argv[1];
		cfg->gain = strtod(argv[2], &end);
		if (*end != '\0' || errno == ERANGE || cfg->gain <= 1.000001f) {
			print_error("Error: invalid value for gain or out of range\n");
			goto _parse_abort;
		}
		cfg->offset_ms = strtoul(argv[3], &end, 10);
		if (*end != '\0' || errno == ERANGE) {
			print_error("Error: invalid value for offset\n");
			goto _parse_abort;
		}
		break;
	default:
		print_error("Error: wrong number of arguments\n");
		goto _parse_abort;
	}

	/* No error found */
	return;

_parse_abort:
	print_usage(argv[0]);
	exit(-4);
}

static void print_usage(const char *app_name)
{
	printf("Usage:\n");
	printf("    %s <original_audio_file> <gain> <offset> [p]\n\n", app_name);
	printf(" If `p` flag is present, then it will attempt to play all three\n");
	printf(" audio samples (original, saturated/clipped, declipped) at the end\n");
	printf(" of execution.\n\n");
	printf("Caveats:\n");
	printf(" - Audio file should be in RAW format (mono, signed 16 bits/sample @ 8 kHz).\n");
	printf(" - Gain must be a real number greater than 1.\n");
	printf(" - The offset should point to a valid position within\n");
	printf("   the audio file.\n");
}