/******************************************************************************
 BROT:
 A simple mandlebrot and julia set program. Written by Neil Robertson.

 Version: 20060913

 Keys:
   R     - Reset
   Q     - Quit
   C     - Reset colour cycling amount
   I     - Toggle info
   1     - Decrease colour cycling amount
   2     - Increase colour cycling amount
   Left  - Decrement julia real
   Right - Increment julia real
   Up    - Increment julia imaginary
   Down  - Decrement julia imaginary

 ******************************************************************************/

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <math.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/keysym.h>

#define WIDTH 640
#define HEIGHT 480
#define NUM_COLOURS 120
#define NUM_UNIQUE_COLOURS 106
#define M_START -2
#define N_START -1.5
#define M_END 1.5
#define N_END 1.5
#define MIN_DIST 5
#define JULIA_REAL 0.233
#define JULIA_IMAG 0.55
#define JULIA_INC 0.001

/******************************** FORWARD DECL *******************************/
void parseCmdLine(int argc, char **argv);
void Xinit();
void init();
void mainloop();
void draw();
void drawBox(int sx, int sy, int ex, int ey);
void printInfo();
void calcNewCords();
int modColour(int col);


/********************************* GLOBALS ***********************************/

Display *display;
char *disp;
Window win;
GC gcw,gcb,gcol[NUM_COLOURS];
int col_cyc_add;
int julia_set;
int keep_aspect;
int show_info;
double box_sx;
double box_sy;
double box_ex;
double box_ey;
double prev_box_sx;
double prev_box_sy;
double prev_box_ex;
double prev_box_ey;
double width;
double height;
double m_start,m_end;
double n_start,n_end;
double aspect_ratio;
double julia_real;
double julia_imag;
double julia_inc;

/******************************** START HERE **********************************/

int main(int argc, char **argv)
{
	parseCmdLine(argc,argv);
	Xinit();
	init();
	mainloop();
	return 0;
}





/*** Parse whats on the command line ***/
void parseCmdLine(int argc, char **argv)
{
	const char *opt[] = 
	{
		"dp",
		"w",
		"h",
		"ms",
		"me",
		"ns",
		"ne",
		"jreal",
		"jimag",
		"jinc",
		"julia",
		"info",
		"ka"
	};
	enum 
	{
		OPT_DP,
		OPT_W,
		OPT_H,
		OPT_MS,
		OPT_ME,
		OPT_NS,
		OPT_NE,
		OPT_JREAL,
		OPT_JIMAG,	
		OPT_JINC,
		OPT_JULIA,
		OPT_INFO,
		OPT_KA,

		OPT_END
	};

	int i,o;
	double val;

	disp = NULL;
	width = WIDTH;
	height = HEIGHT;
	m_start = M_START;
	m_end = M_END;
	n_start = N_START;
	n_end = N_END;
	julia_real = JULIA_REAL;
	julia_imag = JULIA_IMAG;
	julia_inc = JULIA_INC;
	julia_set = 0;
	keep_aspect = 0;
	show_info = 0;

	for(i=1;i < argc;i+=2)
	{
		if (argv[i][0] != '-') goto USAGE;
		for(o=0;o != OPT_END;++o)
			if (!strcasecmp(opt[o],argv[i]+1)) break;

		switch(o)
		{
		case OPT_JULIA:
		case OPT_INFO:
		case OPT_KA:
			--i;  break;

		default:
			if (i == argc-1) goto USAGE;
			val = atof(argv[i+1]);
		}


		switch(o)
		{
		case OPT_DP:
			disp = argv[i+1];  break;

		case OPT_W:
			if (val < 1)
			{
				puts("ERROR: Width must be > 0.");
				exit(1);
			}
			width = val;
			break;

		case OPT_H:
			if (val < 1)
			{
				puts("ERROR: Height must be > 0.");
				exit(1);
			}
			height = val;
			break;

		case OPT_MS:
			m_start = val;
			break;

		case OPT_ME:
			m_end = val;
			break;

		case OPT_NS:
			n_start = val;
			break;

		case OPT_NE:
			n_end = val;
			break;

		case OPT_JREAL:
			julia_real = val;
			break;

		case OPT_JIMAG:
			julia_imag = val;
			break;

		case OPT_JINC:
			julia_inc = val;
			break;
			
		case OPT_JULIA:
			julia_set = 1;
			break;

		case OPT_INFO:
			show_info = 1;
			break;

		case OPT_KA:
			keep_aspect = 1;
			break;

		default: goto USAGE;
		}
	}
	return;

	USAGE:
	printf("Usage: %s\n"
	       "       -dp    <display>\n"
	       "       -w     <window width>         : Default = %d\n"
	       "       -h     <window height>        : Default = %d\n"
               "       -ms    <M start value>        : Default = %f\n"
               "       -me    <M end value>          : Default = %f\n"
               "       -ns    <N start value>        : Default = %f\n"
               "       -ne    <N end value>          : Default = %f\n"
	       "       -jreal <julia set real>       : Default = %f\n"
	       "       -jimag <julia set imaginary>  : Default = %f\n"
	       "       -jinc  <real/imag increments> : Default = %f\n"
	       "       -julia                        : Use julia set algorithm\n"
	       "       -info                         : Print image info to stdout\n"
	       "       -ka                           : Keep window aspect ratio when selecting\n"
	       "                                      next area.\n"
	       "All arguments are optional. The -jreal and -jimag arguments have no effect\nif the -julia argument is not specified.\n",
		argv[0],
		WIDTH,
		HEIGHT,
		(float)M_START,
		M_END,
		N_START,
		N_END,
		JULIA_REAL,
		JULIA_IMAG,
		JULIA_INC);
	exit(1);
}




/*** Set up X ***/
void Xinit()
{
	Window rootwin;
	XGCValues gcvals;
	Colormap cmap;
	XColor col;
	XColor unused;
	XTextProperty title;
	int screen,white,black;
	int cnt,stage;
	unsigned char r,g,b;
	char colstr[5];
	char *tstr = "MANDLEBROT & JULIA SETS";

	if (!(display = XOpenDisplay(disp)))
	{
		printf("ERROR: Can't connect to display %s\n",XDisplayName(disp));
		exit(1);
	}

	screen = DefaultScreen(display);
	rootwin = RootWindow(display,screen);
	cmap = DefaultColormap(display,screen);
	white = WhitePixel(display,screen);
	black = BlackPixel(display,screen);

	win = XCreateSimpleWindow(
		display,rootwin,0,0,(int)width,(int)height,0,white,black);

	/* Set title */
	XStringListToTextProperty(&tstr,1,&title);
	XSetWMProperties(display,win,&title,NULL,NULL,0,NULL,NULL,NULL);

	/* Create black & white GCs */
	gcvals.foreground = black;
	gcb = XCreateGC(display,win,GCForeground,&gcvals);

	gcvals.foreground = white;
	gcvals.function = GXxor;
	gcw = XCreateGC(display,win,GCForeground | GCFunction,&gcvals);

	/* Create colour GCs */
	r = 0;
	g = 0;
	b = 1;
	stage = 0;

	for(cnt=0;cnt < NUM_COLOURS;++cnt) 
	{
		sprintf(colstr,"#%01X%01X%01X",r,g,b);
		if (!XAllocNamedColor(display,cmap,colstr,&col,&unused)) 
		{
			printf("WARNING: Can't allocate colour %s\n",colstr);
			col.pixel = white;
		}
		gcvals.foreground = col.pixel;
		gcol[cnt] = XCreateGC(display,win,GCForeground,&gcvals);

		switch(stage)
		{
		case 0:
			if (++b == 0xF) stage = 1;
			break;

		case 1:
			if (++g == 0xF) stage = 2;
			break;

		case 2:
			if (!--b) stage = 3;
			break;

		case 3:
			if (++r == 0xF) stage = 4;
			break;

		case 4:
			if (!--g) stage = 5;
			break;

		case 5:
			if (++b == 0xF) stage = 6;
			break;

		case 6:
			if (++g == 0xF) stage = 7;
			break;

		case 7:
			/* Colours here are used purely for smooth colour 
			   return to the start colour */
			--r;
			--g;
			--b;
		}
	}

	XSelectInput(display,win,
		ExposureMask | 
		KeyPressMask | 
		ButtonPressMask | 
		ButtonReleaseMask | 
		PointerMotionMask |
		StructureNotifyMask);

	XMapWindow(display,win);
}





/*** Initialise non X stuff ***/
void init()
{
	aspect_ratio = width / height;

	prev_box_sx = -1;
	prev_box_sy = -1;
	prev_box_ex = -1;
	prev_box_ey = -1;

	col_cyc_add = 0;
}




/*** Run the thing ***/
void mainloop()
{
	Window rw,cw;
	XEvent event;
	KeySym ksym;
	char key;
	int rx,ry,wx,wy;
	int button_pressed;
	u_int unused;

	for(button_pressed=0;;)
	{
		XNextEvent(display,&event);

		switch(event.type)
		{
		case Expose:
			draw();  break;

		case ConfigureNotify:
			if (event.xconfigure.width != (int)width ||
			    event.xconfigure.height != (int)height)
			{
				width = event.xconfigure.width;
				height = event.xconfigure.height;
				aspect_ratio = width / height;

				draw();
			}
			break;

		case KeyPress:
			XLookupString(&event.xkey,&key,1,&ksym,NULL);
			switch(ksym)
			{
			case XK_c:
			case XK_C:
				col_cyc_add = 0;
				break;

			case XK_i:
			case XK_I:
				show_info = !show_info;
				printInfo();
				break;

			case XK_q:
			case XK_Q:
				exit(0);

			case XK_r:
			case XK_R:
				m_start = M_START;
				n_start = N_START;
				m_end = M_END;
				n_end = N_END;
				break;

			case XK_1:
				col_cyc_add--;
				if (col_cyc_add < 0)
					col_cyc_add = NUM_COLOURS - 1;
				break;

			case XK_2:
				col_cyc_add = (col_cyc_add + 1) % NUM_COLOURS;
				break;

			case XK_Left:
				julia_real -= julia_inc;
				draw();
				break;

			case XK_Right:
				julia_real += julia_inc;
				draw();
				break;

			case XK_Up:
				julia_imag += julia_inc;
				draw();
				break;

			case XK_Down:
				julia_imag -= julia_inc;
				draw();
				break;

			}
			draw();
			break;

		case ButtonPress:
			if (event.xbutton.button == 1) 
			{
				button_pressed = 1;
				XQueryPointer(
					display,win,&rw,&cw,
					&rx,&ry,&wx,&wy,&unused);
				box_ex = box_sx = (double)wx;
				box_ey = box_sy = (double)wy;
			}
			break;

		case ButtonRelease:
			if (event.xbutton.button == 1) 
			{
				button_pressed = 0;
				drawBox(
					(int)prev_box_sx,(int)prev_box_sy,
					(int)prev_box_ex,(int)prev_box_ey);
				prev_box_sx = -1;
				prev_box_sy = -1;
				prev_box_ex = -1;
				prev_box_ey = -1;

				if (fabs(box_ex - box_sx) >= MIN_DIST &&
				    fabs(box_ey - box_sy) >= MIN_DIST)
				{
					calcNewCords();
					draw();
				}
			}
			break;

		case MotionNotify:
			if (!button_pressed) break;

			box_ex = event.xmotion.x;
			if (keep_aspect)
				box_ey = box_sy + (box_ex - box_sx) / aspect_ratio;
			else
				box_ey = event.xmotion.y;

			drawBox(
				(int)prev_box_sx,(int)prev_box_sy,
				(int)prev_box_ex,(int)prev_box_ey);

			drawBox((int)box_sx,(int)box_sy,(int)box_ex,(int)box_ey);

			prev_box_sx = box_sx;
			prev_box_sy = box_sy;
			prev_box_ex = box_ex;
			prev_box_ey = box_ey;
			break;
		}
	}
}




/*** Draw the set ***/
void draw()
{
	double x,y;
	double m,n,m1,n1,m2,n2;
	double m_add,n_add;
	double real, imag;
	int i;

	m_add = (m_end - m_start) / width;
	n_add = (n_end - n_start) / height;

	for(x=0,m=m_start;x < width;++x,m+=m_add)
	{
		for(y=0,n=n_start;y < height;++y,n+=n_add)
		{
			m1 = m;
			n1 = n;
			real = (julia_set ? julia_real : m);
			imag = (julia_set ? julia_imag : n);

			for(i=0;i < NUM_UNIQUE_COLOURS && m1*m1 + n1*n1 < 4;++i)
			{
				m2 = m1 * m1 - n1 * n1 + real;
				n2 = 2 * m1 * n1 + imag;

				m1 = m2;
				n1 = n2;
			}
			
			if (i == NUM_UNIQUE_COLOURS)
				XDrawPoint(display,win,gcb,(int)x,(int)y);
			else
				XDrawPoint(display,win,gcol[modColour(i)],(int)x,(int)y);
		}
	}
	XFlush(display);

	if (show_info) printInfo();
}




/*** Draw the bounding box for zooming in ***/
void drawBox(int sx, int sy, int ex, int ey)
{
	XDrawLine(display,win,gcw,sx,sy,ex,sy);
	XDrawLine(display,win,gcw,sx,sy,sx,ey);
	XDrawLine(display,win,gcw,ex,sy,ex,ey);
	XDrawLine(display,win,gcw,sx,ey,ex,ey);
}




/*** Print the complex number plane info ***/
void printInfo()
{
	if (julia_set)
		printf("\rM: %.4f -> %.4f  N: %.4f -> %.4f   Real: %.4f  Imag: %.4f",
			m_start,m_end,n_start,n_end,julia_real,julia_imag);
	else
		printf("\rM: %.4f -> %.4f   N: %.4f -> %.4f",
			m_start,m_end,n_start,n_end);
	fflush(stdout);
}




/*** Calculate new M & N complex plain co-ords from box start and end
     positions ***/
void calcNewCords()
{
	double tmp,dx,dy;

	if (box_ex < box_sx)
	{
		tmp = box_ex;
		box_ex = box_sx;
		box_sx = tmp;
	}
	if (box_ey < box_sy)
	{
		tmp = box_ey;
		box_ey = box_sy;
		box_sy = tmp;
	}

	dx = m_end - m_start;
	dy = n_end - n_start;

	if (keep_aspect) 
		box_ey = box_sy + (box_ex - box_sx) / aspect_ratio;

	m_start = m_start + dx * box_sx / width;
	n_start = n_start + dy * box_sy / height;

	m_end = m_start + dx * (box_ex - box_sx) / width;
	n_end = n_start + dy * (box_ey - box_sy) / height;
}




/*** Modify the colour ***/
int modColour(int col)
{
	if (col > NUM_UNIQUE_COLOURS) col = NUM_UNIQUE_COLOURS - 1;
	return (col + col_cyc_add) % NUM_COLOURS;
}