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

#include <gsl/gsl_spline.h>

#include "pnm.h"

typedef struct {
	int n;
	unsigned short value[1];
} table_t;

static int
pam_transform(FILE *in_fp, FILE *out_fp,
	const struct pnm *in_pnm, table_t **tables);

static table_t*
buildCurve(const char *ctrl_points, int resolution, int maxval);

#define RED 0x1
#define GREEN 0x2
#define BLUE 0x4

int
main(int argc, char **argv) {
	int c, i;
	char *curve = NULL;
	struct pnm in_pnm;
	int channels = 0;
	table_t *table, *gTable[3];

	while ((c = getopt(argc, argv, "rgbc:h")) != EOF) {
		switch (c) {  
			case 'r':
				channels |= RED;
				break;
			case 'g':
				channels |= GREEN;
				break;
			case 'b':
				channels |= BLUE;
				break;
			case 'c':
				curve = optarg;
				break;

			default:
				break;
		}
	}

	if (readPnmHeader(stdin, &in_pnm) != 0) {
		fprintf(stderr, "failed to read PNM file\n");
		exit(1);
	}

	if (in_pnm.maxval != 255 && in_pnm.maxval != 65535) {
		fprintf(stderr, "unsupported PNM maxval %d\n", in_pnm.maxval);
		exit(1);
	}

	if (curve) {
		if (channels == 0) channels = RED | GREEN | BLUE;

		table = buildCurve(curve, in_pnm.maxval + 1, in_pnm.maxval);

		if (!table) {
			fprintf(stderr, "could not build table.\n");
			exit(1);
		}
	
		gTable[0] = channels & RED ? table : NULL;	
		gTable[1] = channels & GREEN ? table : NULL;	
		gTable[2] = channels & BLUE ? table : NULL;	
	}

	writePnmHeader(stdout, &in_pnm);

	pam_transform(stdin, stdout, &in_pnm, gTable);

	return 0;
}

static int
pam_transform(FILE *in_fp, FILE *out_fp,
	const struct pnm *in_pnm, table_t **tables) {
	int i, c, row, col;
	int nbytes = in_pnm->maxval == 65535?2:1;
	unsigned char *buf = malloc(in_pnm->width * nbytes * 3);

	for (c = 0; c < 3; c++) {
		if (tables[c] && tables[c]->n <= in_pnm->maxval) {
			fprintf(stderr, "table %d too small (%d)\n", c, tables[c]->n);
			return 1;
		}
	}

	for (row = 0; row < in_pnm->height; row++) {
		fread(buf, in_pnm->width, nbytes * 3, in_fp);

		for (i = 0; i < in_pnm->width; i++) {
			for (c = 0; c < 3; c++) {
				if (tables[c]) {
					int val = buf[(i * 3 + c) * nbytes];
					if (nbytes == 2) {
						val = (val << 8) | buf[(i * 3 + c) * nbytes + 1];
					}

					val = tables[c]->value[val];

					if (nbytes == 1) {
						buf[i * 3 + c] = (unsigned char) val;
					} else if (nbytes == 2) {
						buf[(i * 3 + c) * 2] = (unsigned char) (val >> 8);
						buf[(i * 3 + c) * 2 + 1] = (unsigned char) val;
					}
				}
			}
		}

		fwrite(buf, in_pnm->width, nbytes * 3, out_fp);
	}

	free(buf);
	return 0;
}

#define MAX_CTRL 256

static table_t*
buildCurve(const char *ctrl_points, int resolution, int maxval) {
	table_t *table;
	char *pstr, *buf = strdup(ctrl_points);
	int i, n = 0;
	double X[MAX_CTRL], Y[MAX_CTRL];
	gsl_interp_accel *acc;
	gsl_spline *spline;

	while (pstr = strsep(&buf, ",")) {
		if (n >= MAX_CTRL) {
			fprintf(stderr, "maximum number of control points (%d) reached.\n",
				MAX_CTRL);
			break;
		}

		if (sscanf(pstr, "%lf:%lf", &X[n], &Y[n]) != 2 ||
			X[n] < 0.0 || X[n] > 1.0 || Y[n] < 0.0 || Y[n] > 1.0) {
			fprintf(stderr, "could not parse control point %s.\n", pstr);
			free(buf);
			return NULL;
		}

		n++;
	}

	free(buf);

	if (n < 3) {
		fprintf(stderr, "not enough control point specified.\n");
		return NULL;
	}

	table = (table_t*) malloc(sizeof(table_t) + resolution * sizeof(short));
	table->n = resolution;

	acc = gsl_interp_accel_alloc();
	spline = gsl_spline_alloc(gsl_interp_cspline, n);
	gsl_spline_init (spline, X, Y, n);

	for (i = 0; i < resolution; i++) {
		double _x = (double) i / (resolution - 1);
		double _y = gsl_spline_eval(spline, _x, acc);
        double val = _y * maxval + 0.5;
        if (val > maxval) val = (double) maxval;
        if (val < 0) val = 0;

        table->value[i] = (unsigned short) floor(val);
	}

	gsl_spline_free(spline);
	gsl_interp_accel_free(acc);
	
	return table;
}