#include "SDL.h"
#include <math.h>
#include <unistd.h>
#include <getopt.h>
#include <ctype.h>
#include <zlib.h>

/* Algorithmic variables */
int maxiters = 256;
#define BAILOUT 2.0
int xres = 600;
int yres = 600;
double xmin = -2, ymin = -2, xmax = 2, ymax = 2;

/* Display variables */
SDL_Surface* screen;
SDL_Color palette[256],origpalette[256];
int ppm = 0;
int sdl = 1;
int cycling = 0;

int iterate(double x, double y)
{
	int i;
	double cplxr, cplxi;
	cplxr = y;
	cplxi = x;
	i = 0;
	while ((i < maxiters) && (hypot(cplxr, cplxi) < BAILOUT))
	{
		double a=cplxr, b=cplxi, c=cplxr, d=cplxi;
		i++;
		cplxr = a*c-b*d+y;
		cplxi = a*d+b*c+x;
	}
	return i;
}

void set_palette()
{
	if (sdl)
		SDL_SetColors(screen, palette, 0, 256);
}

void cycle_colors()
{
	SDL_Color x;
	int i;
	x = palette[0];
	for (i=0; i<255; i++)
		palette[i] = palette[i+1];
	palette[255] = x;
	set_palette();
}

Uint32 cycle_colors_cbk(Uint32 interval)
{
	cycle_colors();
	return interval;
}

void init_palette()
{
	int i;
	for (i=0; i<256; i++)
	{
		palette[i].r = sin(((double)i)/256.0*2.0*M_PI*5.0f)*256.0;
		palette[i].g = sin(((double)(i+64))/256.0*2.0*M_PI*4.0f)*256.0;
		palette[i].b = sin(((double)(i+128))/256.0*2.0*M_PI*3.0f)*256.0;
		origpalette[i] = palette[i];
	}
}

int videoflags = SDL_SWSURFACE | SDL_RESIZABLE;

void start_sdl()
{
	if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_TIMER) < 0)
	{
		printf("SDL init failed: %s\n", SDL_GetError());
		exit(1);
	}
	atexit(SDL_Quit);
	
	screen = SDL_SetVideoMode(xres, yres, 8, videoflags);
	if (!screen)
	{
		printf("SDL video init failed: %s\n", SDL_GetError());
		exit(1);
	}
	
	SDL_WM_SetCaption("Mandelbrot Set", "Mandelbrot");
}

double pxltofltx(int x)
{
	return ((double)x)*((double)(xmax-xmin))/((double)xres)+xmin;
}
double pxltoflty(int y)
{
	return ((double)y)*((double)(ymax-ymin))/((double)yres)+ymin;
}

void render()
{
	int x,y;
	Uint8* buffer;
	if (ppm)
		printf("P6\n%d %d\n255\n", xres, yres);
	if (sdl)
	{
		cycling = 0;
		SDL_WM_SetCaption("Mandelbrot Set [working]", "Mandelbrot");
		SDL_SetTimer(0, NULL);
		SDL_LockSurface(screen);
		buffer = (Uint8*)screen->pixels;
	}
	for (y=0;y<yres;y++)
	{
		for (x=0;x<xres;x++)
		{
			int val = iterate(pxltofltx(x), pxltoflty(y));
			if (ppm)
				printf("%c%c%c", palette[val].r, palette[val].g, palette[val].b);
			if (sdl)
				buffer[x] = val;
		}
		if (sdl)
			buffer += screen->pitch;
		if (!(y%10))
		{
			if (sdl)
			{
				for (x=0;x<xres;x++)
					buffer[x]=0;
				SDL_UpdateRect(screen,0,0,0,0);
			} else 
				fprintf(stderr, "%d\n", y);
		}
	}
	if (sdl)
	{
		SDL_UpdateRect(screen, 0, 0, 0, 0);
		SDL_UnlockSurface(screen);
		SDL_WM_SetCaption("Mandelbrot Set", "Mandelbrot");
	}
}

int load_palette(char* name)
{
	gzFile file;
	char c[1024];
	int idx;
	file = gzopen(name, "rb");
	if (file == NULL)
	{
		perror("gzopen");
		return 0;
	}
	for (idx = 0; idx < 256; idx++)
	{
		char* cidx = c;
		int accum;
		int element;
		gzgets(file, c, 1024);
		
		if ((c[0] == '\r') || (c[0] == '\n'))
		{
			idx--;
			continue;
		}
		if (c[0] == 0x1A /* EOF */)
		{
			gzclose(file);
			return 1;
		}
		for (element = 0; element < 3; element++)
		{
			accum = 0;
			while (isblank(*cidx))
				cidx++;
			if (!isdigit(*cidx))
			{
				fprintf(stderr, "%s: palette file format fault on color %d -- expected digit, got `%c'\n", name, idx, *cidx);
				gzclose(file);
				return;
			}
			while (isdigit(*cidx))
			{
				accum *= 10;
				accum += *cidx - '0';
				cidx++;
			}
			switch(element)
			{
			case 0:	origpalette[idx].r = palette[idx].r = accum; break;
			case 1: origpalette[idx].g = palette[idx].g = accum; break;
			case 2: origpalette[idx].b = palette[idx].b = accum; break;
			default: abort();
			}
		}
	}
	gzclose(file);
	set_palette();
	return 1;
}

void usage(char* me)
{
	printf(
		"Usage: %s [OPTION]...\n"
		"Renders the Mandelbrot set.\n"
		"\n"
		"Mandatory arguments to long options are mandatory for short options too.\n"
		"  -x, --width=PIXELS           sets the output width\n"
		"  -y, --height=PIXELS          sets the output height\n"
		"  -m, --maxiters=ITERATIONS    sets the maximum number of iterations\n"
		"  -s, --sdl                    enables SDL as an output device (default)\n"
		"  -S, --no-sdl                 disables SDL as an output device (implies -q)\n"
		"  -p, --ppm                    enables .ppm to stdout as an output device (implies -q)\n"
		"  -P, --no-ppm                 disables .ppm as an output device (default)\n"
		"  -q, --quit-immediately       exits immediately after rendering\n"
		"  -l, --location=XMIN,XMAX,YMIN,YMAX\n"
		"                               sets the starting location\n"
		"  -t, --palette=FILENAME       sets the palette\n"
		"      --help                   shows this help\n"
		"\n"
		"Written by Joshua Wise <joshua@joshuawise.com>.\n"
		, me
		);
	exit(2);
}

struct option longopts[] = {
	{"x", required_argument, NULL, 'x'},
	{"width", required_argument, NULL, 'x'},
	{"y", required_argument, NULL, 'y'},
	{"height", required_argument, NULL, 'y'},
	{"maxiters", required_argument, NULL, 'm'},
	{"sdl", no_argument, NULL, 's'},
	{"ppm", no_argument, NULL, 'p'},
	{"no-sdl", no_argument, NULL, 'S'},
	{"no-ppm", no_argument, NULL, 'P'},
	{"help", no_argument, NULL, 'h'},
	{"location", required_argument, NULL, 'l'},
	{"palette", required_argument, NULL, 't'},
	{"quit-immediately", no_argument, NULL, 'q'},
	{0,0,0,0}
};

int main(int argc, char** argv)
{
        SDL_Event event;
        int startx, starty;
        int buttondown = 0;
        int dragged = 0;
        int arg;
        int custompalette = 0;
        int quitimmediately = 0;
        double yctr, ysz;
        
        while ((arg = getopt_long(argc, argv, "x:y:m:l:t:sSpPq", longopts, NULL)) != -1)
        {
        	switch(arg)
        	{
		case 'x':
			xres = atoi(optarg);
			if (!xres)
			{
				printf("%s: option `%c' requires an integer argument\n", argv[0], arg);
				printf("Try `%s --help' for more information.\n", argv[0]);
				exit(2);
			}
			break;
		case 'y':
			yres = atoi(optarg);
			if (!yres)
			{
				printf("%s: option `%c' requires an integer argument\n", argv[0], arg);
				printf("Try `%s --help' for more information.\n", argv[0]);
				exit(2);
			}
			break;
		case 'm':
			maxiters = atoi(optarg);
			if (!maxiters)
			{
				printf("%s: option `%c' requires an integer argument\n", argv[0], arg);
				printf("Try `%s --help' for more information.\n", argv[0]);
				exit(2);
			}
			break;
		case 's':
			sdl = 1;
			break;
		case 'p':
			ppm = 1;
			quitimmediately = 1;
			break;
		case 'S':
			sdl = 0;
			break;
		case 'P':
			ppm = 0;
			break;
		case 'h':
			usage(argv[0]);
			break;
		case 'l':
			{
			float f1,f2,f3,f4;
				if (sscanf(optarg, "%f,%f,%f,%f", &f1, &f2, &f3, &f4) != 4)
				{
					printf("%s: incorrect syntax for option `%c'\n", argv[0], arg);
					printf("Try %s --help' for more information.\n", argv[0]);
					exit(2);
				}
				xmin = f1; xmax = f2; ymin = f3; ymax = f4;
			}
			break;
		case 't':
			custompalette = load_palette(optarg);
			break;
		case 'q':
			quitimmediately = 1;
			break;
		case '?':
		case ':':
			printf("Try `%s --help' for more information.\n", argv[0]);
			exit(2);
			break;
		default:
			printf("Oh God I am not good with computers how did this get here?\n");
			exit(127);
        	}
        }
        
	if (sdl)
		start_sdl();
	yctr = ymin + (ymax - ymin) / 2.0;
	// get y in proportion
	ysz = (xmax-xmin)*((double)yres)/((double)xres);
	ymax = yctr+ysz/2; ymin = yctr-ysz/2;

	if (!custompalette)
		init_palette();
	set_palette();
	render();
	if (quitimmediately)
		exit(0);
	while ( SDL_WaitEvent(&event) ) {
                switch (event.type) {
                	case SDL_KEYDOWN:
                		if (event.key.keysym.sym == SDLK_c)
                		{
                			if (!cycling)
                			{
                				SDL_WM_SetCaption("Mandelbrot Set [color cycling]", "Mandelbrot");
                				SDL_SetTimer(50, cycle_colors_cbk);
					} else {
						SDL_WM_SetCaption("Mandelbrot Set", "Mandelbrot");
						SDL_SetTimer(0, NULL);
					}
                			cycling = !cycling;
                		} else if (event.key.keysym.sym == SDLK_r) {
                			int i;
                			for (i=0;i<256;i++)
                			{
                				palette[i].r = origpalette[i].r;
                				palette[i].g = origpalette[i].g;
                				palette[i].b = origpalette[i].b;
					}
					set_palette();
                		} else if (event.key.keysym.sym == SDLK_ESCAPE ||
					   event.key.keysym.sym == SDLK_q) {
					exit(0);
				} else if (event.key.keysym.sym == SDLK_p) {
					ppm = 1;
					render();
					ppm = 0;
				} else if (event.key.keysym.sym == SDLK_RETURN) {
					render();
				} else if (event.key.keysym.sym == SDLK_i) {
					printf("mandelbrot by Joshua Wise\n");
					printf("2005-06-19, Little Cayman\n");
					printf("\n");
					printf("Current settings:\n");
					printf("  xmin: %f, xmax: %f\n", xmin, xmax);
					printf("  ymin: %f, ymax: %f\n", ymin, ymax);
					printf("  maxiters: %d, bailout: %f\n", maxiters, BAILOUT);
					printf("  xres: %d, yres: %d\n", xres, yres);
					printf("Recreate with:\n");
					printf("  %s -x%d -y%d -l%f,%f,%f,%f -m%d\n", argv[0], xres,yres,xmin,xmax,ymin,ymax,maxiters);
				}
                		break;
			case SDL_MOUSEBUTTONDOWN:
				if (event.button.button == 1)
					buttondown = 1;
				startx = event.button.x;
				starty = event.button.y;
				break;
			case SDL_MOUSEMOTION:
				if (buttondown)
					dragged = 1;
				break;
			case SDL_MOUSEBUTTONUP:
				if (event.button.button == 1)
					buttondown = 0;
				if (!dragged)
				{
					if (event.button.button == 1 || event.button.button == 3)
					{
						double xsz = xmax-xmin, ysz = ymax-ymin;
						double x,y;
						double scalar;
						if (event.button.button == 1)
							scalar = 1.0/2.0;
						else
							scalar = 2.0/1.0;
						xsz *= scalar;
						ysz *= scalar;
						x = pxltofltx(event.button.x);
						y = pxltoflty(event.button.y);
						xmax = x + xsz / 2;
						xmin = x - xsz / 2;
						ymax = y + ysz / 2;
						ymin = y - ysz / 2;
						render();
					} else if (event.button.button == 2) {
						xmax = ymax = 2;
						xmin = ymin = -2;
						render();
					}
					dragged = 0;
					break;
				} else {
					double newmaxx, newminx, newmaxy, newminy;
					if (event.button.x > startx)
					{
						newmaxx = pxltofltx(event.button.x);
						newminx = pxltofltx(startx);
					} else {
						newminx = pxltofltx(event.button.x);
						newmaxx = pxltofltx(startx);
					}
					if (event.button.y > starty)
					{
						newmaxy = pxltoflty(event.button.y);
						newminy = pxltoflty(starty);
					} else {
						newminy = pxltoflty(event.button.y);
						newmaxy = pxltoflty(starty);
					}
					yctr = newminy + (newmaxy - newminy) / 2.0;
					// get y in proportion
					ysz = (newmaxx-newminx)*((double)yres)/((double)xres);
					ymax = yctr+ysz/2; ymin = yctr-ysz/2;
					xmax = newmaxx; xmin = newminx;
					render();
					dragged = 0;
					break;
				}
			case SDL_VIDEORESIZE:
			{
				screen = SDL_SetVideoMode(event.resize.w, event.resize.h,
							  screen->format->BitsPerPixel,
							  videoflags);
				if (!screen)
				{
					fprintf(stderr, "Couldn't resize screen\n");
					exit(1);
				}
				SDL_SetColors(screen, palette, 0, 256);
				xres = event.resize.w;
				yres = event.resize.h;
				yctr = ymin + (ymax - ymin) / 2.0;
				// get y in proportion
				ysz = (xmax-xmin)*((double)yres)/((double)xres);
				ymax = yctr+ysz/2; ymin = yctr-ysz/2;
				render();
				break;
			}
			case SDL_QUIT:
				return;
		}
	}
	return 1;
}