/* Animated frequency filter
 *
 * Dylan Simon <https://dylex.net:9947/src/>
 */

#include <math.h>
#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <complex.h>
#include <stdbool.h>
#include <sys/time.h>

#include <fftw3.h>
#include <gd.h>

#include <X11/Intrinsic.h>
#include <X11/IntrinsicP.h>
#include <X11/CoreP.h>
#include <X11/Shell.h>
#include <X11/StringDefs.h>

#define MAX(X, Y) ({ typeof(X) _x = X; typeof(Y) _y = Y; _y > _x ? _y : _x; })
#define MIN(X, Y) ({ typeof(X) _x = X; typeof(Y) _y = Y; _y < _x ? _y : _x; })

enum color {
	COLOR_R = 1,
	COLOR_G = 2,
	COLOR_B = 4,
	COLOR_W = COLOR_R | COLOR_G | COLOR_B
};

struct ffid {
	unsigned iw, ih;
	unsigned fw, fh;
	bool rgb;
	unsigned ni;
	double complex **fi;
	double complex *fv;
	double *iv;
	fftw_plan ft, it;

	XImage *xi;
};
static inline size_t fsize(struct ffid *d) { return d->fw*d->fh*sizeof(*d->fv); }
static inline size_t isize(struct ffid *d) { return d->iw*d->ih*sizeof(*d->iv); }

static void init_ffi(struct ffid *d, unsigned w, unsigned h, unsigned n, bool rgb)
{
	d->iw = w;
	d->ih = h;

	d->fw = d->iw/2+1;
	d->fh = d->ih;

	d->iv = malloc(isize(d));
	d->fv = malloc(fsize(d));

	d->rgb = rgb;

	d->ni = n;
	if (d->rgb) n *= 3;
	d->fi = malloc(n*sizeof(*d->fi));
	unsigned i;
	for (i = 0; i < n; i ++)
		d->fi[i] = malloc(fsize(d));

	d->ft = fftw_plan_dft_r2c_2d(d->ih, d->iw, d->iv, d->fv, FFTW_ESTIMATE);
	d->it = fftw_plan_dft_c2r_2d(d->ih, d->iw, d->fv, d->iv, FFTW_MEASURE);
}

static void fini_ffi(struct ffid *d)
{
	XDestroyImage(d->xi);

	fftw_destroy_plan(d->it);
	fftw_destroy_plan(d->ft);

	unsigned i;
	for (i = 0; i < d->ni; i ++)
		free(d->fi[i]);
	free(d->fi);
	d->ni = 0;

	free(d->fv);
	free(d->iv);

	d->ih = 0;
	d->iw = 0;
}

static gdImagePtr open_image(const char *file)
{
	FILE *f;
	gdImagePtr img = NULL;

	char *e = strrchr(file, '.');
	if (!e || !*++e)
		return img;

	if (!(f = fopen(file, "r")))
	{
		if (errno != EEXIST)
			fprintf(stderr, "%s: %m\n", file);
		return img;
	}

	if (!strcasecmp(e, "png"))
		img = gdImageCreateFromPng(f);
	else if (!strcasecmp(e, "jpg") || !strcasecmp(e, "jpeg"))
		img = gdImageCreateFromJpeg(f);
	else if (!strcasecmp(e, "gif"))
		img = gdImageCreateFromGif(f);
	if (!img)
		fprintf(stderr, "%s: error loading image\n", file);
	fclose(f);
	return img;
}

static gdImagePtr draw_text(unsigned w, unsigned h, char *text, char *font, double size)
{
	if (!w) w = 256;
	if (!h) h = 256;
	if (!text || !*text)
	{
		return gdImageCreateTrueColor(w, h);
	}
	int b[8];
	double s = size ?: w/strlen(text);
	int x, y;
	while (1)
	{
		char *e = gdImageStringFT(NULL, b, 0, font, s, 0, 0, 0, text);
		if (e)
		{
			fprintf(stderr, "error drawing text: %s\n", e);
			return NULL;
		}
		x = b[2] - b[6];
		y = b[3] - b[7];
		if (size != 0) break;
		if (x > w || y > h)
			s *= 0.99*MIN((double)w/x, (double)h/y);
		else if (x < 0.97*w && y < 0.97*h)
			s *= 0.99*MIN((double)w/x, (double)h/y);
		else
			break;
	}
	gdImagePtr img = gdImageCreateTrueColor(w, h);
	gdImageStringFT(img, b, gdTrueColor(255,255,255), font, s, 0, (w-x)/2-b[6], (h-y)/2-b[7], text);
	return img;
}

static void read_image(struct ffid *d, const gdImagePtr img, enum color color)
{
	unsigned x, y;
	double n = 1./(d->ih*d->iw)/__builtin_popcountl(color);
	for (y = 0; y < d->ih; y ++) for (x = 0; x < d->iw; x ++)
	{
		int c = gdImageGetPixel(img, x, y);
		double v = 0;
		if (color & COLOR_R)
			v += gdImageRed(img, c) - 128;
		if (color & COLOR_G)
			v += gdImageGreen(img, c) - 128;
		if (color & COLOR_B)
			v += gdImageBlue(img, c) - 128;
		d->iv[d->iw*y+x] = n*v;
	}
}

static void load_image(struct ffid *d, unsigned i, const gdImagePtr img)
{
	if (d->rgb)
	{
		unsigned c;
		for (c = 0; c < 3; c ++)
		{
			read_image(d, img, 1 << c);
			fftw_execute(d->ft);
			memcpy(d->fi[3*i+c], d->fv, fsize(d));
		}
	}
	else
	{
		read_image(d, img, COLOR_W);
		fftw_execute(d->ft);
		memcpy(d->fi[i], d->fv, fsize(d));
	}
}

static void draw_image(struct ffid *d, enum color color)
{
	unsigned x, y;
	double m = 2./d->ni;
	for (y = 0; y < d->ih; y ++) for (x = 0; x < d->iw; x ++)
	{
		double v = d->iv[y*d->iw+x];
		int c = v*m;
		if (c > 255) c = 255;
		else if (c < -255) c = -255;

		if (color & COLOR_B) d->xi->data[d->xi->bytes_per_line*y+d->xi->bits_per_pixel/8*x+0] = c < 0 ? -c : 0;
		if (color & COLOR_G) d->xi->data[d->xi->bytes_per_line*y+d->xi->bits_per_pixel/8*x+1] = c > 0 ? c : 0;
		if (color & COLOR_R) d->xi->data[d->xi->bytes_per_line*y+d->xi->bits_per_pixel/8*x+2] = c > 0 ? c : 0;
	}
}

static void add_filter(struct ffid *d, unsigned n, double lo, double hi)
{
	if (lo > 1 || lo < 0) lo -= floor(lo);
	if (hi > 1 || hi < 0) hi -= floor(hi);

	if (lo > hi)
	{
		add_filter(d, n, 0, hi);
		add_filter(d, n, lo, 1);
		return;
	}

	if (lo >= hi || hi <= 0 || lo >= 1)
		return;

	double s = MIN(d->iw, d->ih);
	lo = pow(s/2, lo);
	hi = pow(s/2, hi);

	void keep(unsigned x, unsigned y, double a)
	{
		d->fv[d->fw*y+x] += a*d->fi[n][d->fw*y+x];
		if (y) d->fv[d->fw*(d->fh-y)+x] += a*d->fi[n][d->fw*(d->fh-y)+x];
	}

	unsigned y = MIN(ceil(hi)-1, d->fh/2);
	for (; (signed)y >= 0; y --)
	{
		double x0 = y >= lo ? 0 : MIN(sqrt(lo*lo - y*y), d->fw);
		double x1 = MIN(sqrt(hi*hi - y*y), d->fw);
		double a = modf(x0, &x0);
		double b = modf(x1, &x1);
		unsigned x = x0;
		if (x0 == x1)
		{
			keep(x, y, b-a);
		}
		else
		{
			if (a > 0)
			{
				keep(x,y,1-a);
				x ++;
			}
			for (; x < x1; x++)
			{
				keep(x,y,1);
			}
			if (b > 0)
			{
				keep(x,y,b);
			}
		}
	}
}

static void filter_images(struct ffid *d, float t, unsigned c)
{
	memset(d->fv, 0, sizeof(*d->fv)*d->fw*d->fh);

	double n = d->ni;
	unsigned i;
	for (i = 0; i < d->ni; i ++)
		add_filter(d, (d->rgb?3:1)*i+c, i/n-t, (i+1)/n-t);

	d->fv[0] = 0;
}

static void make_image(struct ffid *d, float t)
{
	if (d->rgb)
	{
		unsigned c;
		for (c = 0; c < 3; c ++)
		{
			filter_images(d, t, c);
			fftw_execute(d->it);
			draw_image(d, 1 << c);
		}
	}
	else
	{
		filter_images(d, t, 0);
		fftw_execute(d->it);
		draw_image(d, COLOR_W);
	}
}

static Display *Disp;
static Widget App;
static Window Win;

static char *Defaults[] = {
	"ffi.width: 0",
	"ffi.height: 0",
	"ffi.font: /usr/lib/X11/fonts/TTF/Vera.ttf",
	"ffi.speed: 0",
	"ffi.size: 0",
	"ffi.rgb: 0", //BOOL
	0
};

static XrmOptionDescRec Options[] = {
	{"-width", ".width", XrmoptionSepArg, 0},
	{"-height", ".height", XrmoptionSepArg, 0},
	{"-font", ".font", XrmoptionSepArg, 0},
	{"-size", ".size", XrmoptionSepArg, 0},
	{"-speed", ".speed", XrmoptionSepArg, 0},
	{"-rgb", ".rgb", XrmoptionSepArg, 0},
};

static char *get_resource(const char *name)
{
	char *val = NULL;
	XrmValue value;
	char *type;
	if (XrmGetResource(XtDatabase(Disp), name, "FFi.Integer", &type, &value))
		val = strndup(value.addr, value.size);
	return val;
}

#define RESOURCE(VAR, NAME, NONE, EXPR) ({ \
		char *VAR = get_resource(NAME); \
		typeof(EXPR) _res = VAR ? EXPR : NONE; \
		if (VAR) free(VAR); \
		_res; \
	})

int main(int argc, char **argv)
{
	XtAppContext ctx;
	App = XtOpenApplication(&ctx, "ffi", Options, sizeof(Options)/sizeof(*Options), &argc, argv, Defaults, topLevelShellWidgetClass, NULL, 0);
	Disp = XtDisplay(App);

	int w = RESOURCE(s, "ffi.width", 0, atoi(s));
	int h = RESOURCE(s, "ffi.height", 0, atoi(s));

	char *font = get_resource("ffi.font");

	double speed = RESOURCE(s, "ffi.speed", 0, strtod(s, NULL));
	double size = RESOURCE(s, "ffi.size", 0, strtod(s, NULL));

	int rgb = RESOURCE(s, "ffi.rgb", 0, atoi(s));

	int n = argc - 1;
	if (n < 2)
	{
		fprintf(stderr, "Usage: %s IMAGE ...\n", argv[0]);
		XtDestroyWidget(App);
		exit(1);
	}

	if (speed == 0) speed = 4*n;

	gdImagePtr imgs[n];

	int i;
	unsigned mw = 0, mh = 0;
	for (i = 0; i < n; i ++)
	{
		char *arg = argv[i+1];
		if (!(imgs[i] = open_image(arg)))
			imgs[i] = draw_text(w ?: mw, h ?: mh, arg, font, size);

		if (!imgs[i])
			exit(1);

		mw = MAX(mw, gdImageSX(imgs[i]));
		mh = MAX(mh, gdImageSY(imgs[i]));
	}
	
	struct ffid ffi;
	init_ffi(&ffi, w ?: mw, h ?: mh, n, rgb);

	for (i = 0; i < n; i ++)
	{
		gdImagePtr img = imgs[i];
		if (gdImageSX(img) != ffi.iw || gdImageSY(img) != ffi.ih)
		{
			gdImagePtr new = gdImageCreateTrueColor(ffi.iw, ffi.ih);
			gdImageCopyResampled(new, img, 0, 0, 0, 0, ffi.iw, ffi.ih, gdImageSX(img), gdImageSY(img));
			gdImageDestroy(img);
			img = new;
		}

		load_image(&ffi, i, img);
		gdImageDestroy(img);
	}
	
	XtVaSetValues(App, XtNwidth, ffi.iw, XtNheight, ffi.ih, NULL);
	XtRealizeWidget(App);
	Win = XtWindow(App);

	XWindowAttributes attr;
	XGetWindowAttributes(Disp, Win, &attr);
	//Wpm = XCreatePixmap(Disp, Win, attr.width, attr.height, attr.depth);
	ffi.xi = XCreateImage(Disp, attr.visual, attr.depth, ZPixmap, 0, NULL, ffi.iw, ffi.ih, 8, 0);
	ffi.xi->data = malloc(ffi.xi->bytes_per_line*ffi.xi->height);

	GC gc = XCreateGC(Disp, Win, 0, NULL);

	struct timeval t0;
	gettimeofday(&t0, NULL);
	unsigned fc = 0;
	double lt = 0;
	do
	{
		struct timeval t;
		gettimeofday(&t, NULL);
		double dt = t.tv_sec - t0.tv_sec + (double)(t.tv_usec - t0.tv_usec)/1000000;
		if (dt > lt + 5)
		{
			printf("%f FPS\n", fc / dt);
			lt = dt;
		}
		dt /= speed;

		make_image(&ffi, dt);
		XPutImage(Disp, Win, gc, ffi.xi, 0, 0, 0, 0, ffi.iw, ffi.ih);

		XSync(Disp, False);
	} while (++ fc);
	sleep(1);

	fini_ffi(&ffi);
	XtDestroyWidget(App);
	XtDestroyApplicationContext(ctx);
	return 0;
}