#include <u.h>
#include <libc.h>
#include <draw.h>
#include <thread.h>
#include <mouse.h>
#include <cursor.h>
#include <keyboard.h>

enum { ncolors = 1200 };
 
int newtonpow = 3;
double xmin = -2, xmax = 1;
double ymin = -1, ymax = 1;
double jcx = -0.8;
double jcy = 0.156;
double pcoef = -0.7;
char txtbuf[128];
int accuracy = 200;
int fractal = 0;
Mousectl *mctl;
Keyboardctl *kctl;
Rendez rend;
int stopdraw;
uchar *imagedata;
Image *image;
extern Cursor reading;
uchar colors[3 * (ncolors + 1)];

typedef struct Job Job;
struct Job {
	int y0, y1;
	Channel *done;
};

double
convx(Rectangle *r, int x)
{
	return (xmax - xmin) * (x - r->min.x) / (r->max.x - r->min.x) + xmin;
}

double
convy(Rectangle *r, int y)
{
	return (ymax - ymin) * (r->max.y - y) / (r->max.y - r->min.y) + ymin;
}

void
pixel(int x, int y)
{
	draw(screen, Rect(x, y, x + 1, y + 1), display->black, nil, ZP);
}

ulong
iterate(double x, double y)
{
	int i;
	double zx, zy, zx2, zy2, v;
	double cx, cy;
	double px, py, tx, ty;
	double r2;

	switch(fractal){
	
	case 0:
		zx = zy = 0;
		cx = x;
		cy = y;
		break;
	case 1:
		zx = x;
		zy = y;
		cx = jcx;
		cy = jcy;
		break;
	case 2:/* burning ship */
		zx = zy = 0;
		cx = x;
		cy = y;
		break;
	case 3:/* tricorn */
		zx = zy = 0;
		cx = x;
		cy = y;
		break;
	case 4:/* newton */
		zx = x;
		zy = y;
		break;
	case 5:/* phoenix */
		zx = zy = 0;
		px = py = 0;
		cx = x;
		cy = y;
		break;
	}

	zx2 = zx*zx;
	zy2 = zy*zy;
	
	for(i = 0; i < accuracy; i++){
		switch(fractal){
		case 0:
			zy = 2 * zx * zy + cy;
			zx = zx2 - zy2 + cx;
			break;
		case 1:
			zy = 2 * zx * zy + cy;
			zx = zx2 - zy2 + cx;
			break;
		case 2:
			zx = fabs(zx);
			zy = fabs(zy);
			zy = 2 * zx * zy + cy;
			zx = zx2 - zy2 + cx;
			break;
		case 3:
			tx = zy*zx - zy*zy + cx;
			ty = -2*zx*zy + cy;
			zx = tx;
			zy = ty;

			break;

		case 4:
		{
			double a, b;
			double ar, ai;
			double br, bi;
			double dr, di;
			double denom;
			double qr, qi;
			double mag;
			int k, p;
		 
			a = zx;
			b = zy;
		 
			if(fabs(a) < 1e-9 && fabs(b) < 1e-9)
				a = 0.001;
		 
			p = newtonpow;
		 
			for(k = 0; k < 40; k++){
		 
				ar = 1;
				ai = 0;
		 
				for(i = 0; i < p-1; i++){
					double tr = ar*a - ai*b;
					double ti = ar*b + ai*a;
					ar = tr;
					ai = ti;
				}
		 
				dr = p * ar;
				di = p * ai;
		 
				br = ar*a - ai*b;
				bi = ar*b + ai*a;
		  
				br -= 1;
		 
				denom = dr*dr + di*di;
				if(denom < 1e-14)
					break;
		 
				qr = (br*dr + bi*di)/denom;
				qi = (bi*dr - br*di)/denom;
		 
				a -= qr;
				b -= qi;
		 
				mag = br*br + bi*bi;
				if(mag < 1e-12)
					break;
			}
		
			mag = atan2(b, a);
			if(mag < 0)
				mag += 6.2831853;
		 
			return ((int)(mag * 100) + k*8) % ncolors;
		}
		case 5:
			tx = zx*zx - zy*zy + cx + pcoef*px;
			ty = 2*zx*zy + cy + pcoef*py;
			px = zx;
			py = zy;
			zx = tx;
			zy = ty;
			break;

		}
		
		zx2 = zx * zx;
		zy2 = zy * zy;
		r2 = zx*zx + zy*zy;
		if(fractal != 4 && r2 > 4){
			v = i + 1 - log(log(r2)/2)/log(2);
		/*	v = 2 + i - log(log(sqrt(zx2 + zy2)) / log(i + 2)) / 0.69314718; old rendering */
			return (int)(v * 50) % ncolors;
		}
		
		
	}
	return 0;
}

void renderband(void *v)
{
	Job *j = v;
	int x, y;
	uchar *p, *q;
	int w;

	w = Dx(screen->r);
	for(y = j->y0; y < j->y1; y++){
		p = imagedata + 3*w*(y - screen->r.min.y);
		for(x = screen->r.min.x; x < screen->r.max.x; x++){
			q = colors + 3 * iterate(convx(&screen->r, x), convy(&screen->r, y));
			*p++ = q[0];
			*p++ = q[1];
			*p++ = q[2];
		}
	}
	sendul(j->done, 1);
	free(j);
	threadexits(nil);
}

void
drawtxt(void)
{
	snprint(txtbuf, sizeof(txtbuf), "pos: %.4fx+ %.4fx %.4fy %.4fy | Jc: %.3f %+ .3fi | iter: %d | phoenix: %.3f | newton: %d", xmin, xmax, ymin, ymax, jcx, jcy, accuracy, pcoef, newtonpow);
	stringbg(screen, Pt(screen->r.min.x, screen->r.min.y), display->black, ZP, font, txtbuf, display->white, ZP);
}

void
redrawproc(void *)
{
	int x, y;
	uchar *p, *q;

	qlock(rend.l);
	for(;;){
		setcursor(mctl, &reading);
		/*
	old single core stuff:
		p = imagedata;
		for(y = screen->r.min.y; y < screen->r.max.y; y++)
			for(x = screen->r.min.x; x < screen->r.max.x; x++){
				if(stopdraw)
					goto check;
				q = colors + 3 * iterate(convx(&screen->r, x), convy(&screen->r, y));
				*p++ = *q++;
				*p++ = *q++;
				*p++ = *q;
			}
		*/
		int nproc = 16;
		int h, i;
		Channel *done;
		done = chancreate(sizeof(ulong), nproc);
		h = Dy(screen->r) / nproc;
		
		for(i=0; i<nproc; i++){
			Job *j = malloc(sizeof(Job));
			j->y0 = screen->r.min.y + i*h;
			j->y1 = (i == nproc-1) ? screen->r.max.y : j->y0 + h;
			j->done = done;
			proccreate(renderband, j, mainstacksize);
		}
		for(i=0; i<nproc; i++)
			recvul(done);
		
		chanfree(done);

		if(stopdraw)
			goto check;
		lockdisplay(display);
		loadimage(image, image->r, imagedata, Dx(image->r) * Dy(image->r) * 3);
		draw(screen, screen->r, image, nil, screen->r.min);
		drawtxt();
		flushimage(display, 1);
		unlockdisplay(display);
	check:
		stopdraw = 0;
		setcursor(mctl, nil);
		rsleep(&rend);
		stopdraw = 0;
	}
}

void
zoom(void)
{
	Rectangle r;
	double xmin_, xmax_, ymin_, ymax_;

	r = getrect(3, mctl);
	if(r.min.x == 0 && r.min.y == 0 && r.max.x == 0 && r.max.y == 0)
		return;
	xmin_ = convx(&screen->r, r.min.x);
	xmax_ = convx(&screen->r, r.max.x);
	ymin_ = convy(&screen->r, r.max.y);
	ymax_ = convy(&screen->r, r.min.y);
	stopdraw = 1;
	qlock(rend.l);
	xmin = xmin_;
	xmax = xmax_;
	ymin = ymin_;
	ymax = ymax_;
	rwakeup(&rend);
	qunlock(rend.l);
}

char *menus[] = {
	"zoom in",
	"iter+",
	"iter-",
	"julia+",
	"julia-",
	"reset",
	"quit",
	nil,
};

Menu menu = {
	.item = menus
};

char *menus2[] = {
	"mandelbrot",
	"julia",
	"burning ship",
	"tricorn",
	"newton",
	"phoenix",
	nil,
};

Menu menu2 = {
	.item = menus2
};

void
resizethread(void *)
{
	ulong l;

	for(;;){
		if(recv(mctl->resizec, &l) < 1)
			continue;
		stopdraw = 1;
		qlock(rend.l);
		if(getwindow(display, Refnone) < 0)
			sysfatal("getwindow: %r");
		freeimage(image);
		free(imagedata);
		image = allocimage(display, screen->r, RGB24, 0, DBlack);
		imagedata = malloc(3 * Dx(screen->r) * Dy(screen->r));
		rwakeup(&rend);
		qunlock(rend.l);
	}
}

void
initcolors(void)
{
	uchar *p;
	int h, x;

	for(p = colors + 3, h = 0; p < colors + nelem(colors); h++){
		x = 0xFF - abs(h % (ncolors / 3) - ncolors / 6) * 0xFF / (ncolors / 6);
		if(h < ncolors/6){
			*p++ = 0xFF;
			*p++ = x;
			*p++ = 0;
		}else if(h < ncolors/3){
			*p++ = x;
			*p++ = 0xFF;
			*p++ = 0;
		}else if(h < ncolors/2){
			*p++ = 0;
			*p++ = 0xFF;
			*p++ = x;
		}else if(h < 2*ncolors/3){
			*p++ = 0;
			*p++ = x;
			*p++ = 0xFF;
		}else if(h < 5*ncolors/6){
			*p++ = x;
			*p++ = 0;
			*p++ = 0xFF;
		}else{
			*p++ = 0xFF;
			*p++ = 0;
			*p++ = x;
		}
	}
}

void
handlekey(Rune r)
{
	stopdraw = 1;
        qlock(rend.l);
 
        switch(r){

        case 'q':
		threadexitsall(nil);
		break;
        case 'r':
		xmin = -2;
		xmax = 1;
		ymin = -1;
		ymax = 1;
		break;
	case 'm':
		pcoef += 0.01;
		break;
		case 'n':
		pcoef -= 0.01;
		break;
        case 'w':
		jcx -= 0.01;
		break;
        case 'e':
		jcx += 0.01;
		break;
        case 's':
		jcy -= 0.01;
 		break;
        case 'd':
		jcy += 0.01;
		break; 
        case 'x':
		newtonpow--;
		break;
        case 'c':
		newtonpow++;
		break;
	case 'u':
		accuracy -= 10;
		break;
	case 'i':
		accuracy += 10;
		break; 
	}
	rwakeup(&rend);
	qunlock(rend.l);
}

void
threadmain()
{
	static QLock ql;
	int rc;
	Rune r;

	if(initdraw(nil, nil, "fractal") < 0)
		sysfatal("initdraw: %r");
	display->locking = 1;
	unlockdisplay(display);
	if((mctl = initmouse(nil, screen)) == nil)
		sysfatal("initmouse: %r");
	 if((kctl = initkeyboard(nil)) == nil)
         sysfatal("initkeyboard: %r");
	rend.l = &ql;
	image = allocimage(display, screen->r, RGB24, 0, DBlack);
	imagedata = malloc(3 * Dx(screen->r) * Dy(screen->r));
	initcolors();
	proccreate(redrawproc, nil, mainstacksize);
	threadcreate(resizethread, nil, mainstacksize);
	for(;;){
		if(nbrecv(kctl->c, &r))
        		handlekey(r);
		readmouse(mctl);
		if(mctl->buttons & 3){
			lockdisplay(display);
			rc = menuhit(2, mctl, &menu2, nil);
			unlockdisplay(display);
			switch(rc){
			case 0:
				stopdraw = 1;
				qlock(rend.l);
				fractal = 0;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 1:
				stopdraw = 1;
				qlock(rend.l);
				fractal = 1;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 2:
				stopdraw = 1;
				qlock(rend.l);
				xmin = -2.2;
				xmax = 1.2;
				ymin = -2;
				ymax = 0.5;
				fractal = 2;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 3:
				stopdraw = 1;
				qlock(rend.l);
				xmin = -2.5;
				xmax = 2.5;
				ymin = -2;
				ymax = 2;
				fractal = 3;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 4:
				stopdraw = 1;
				qlock(rend.l);
				fractal = 4;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 5:
				stopdraw = 1;
				qlock(rend.l);
				fractal = 5;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			}
		}
		if(mctl->buttons & 4){
			lockdisplay(display);
			rc = menuhit(3, mctl, &menu, nil);
			unlockdisplay(display);
			switch(rc){
			case 0:
				zoom();
				break;
			case 1:
				stopdraw = 1;
				qlock(rend.l);
				accuracy += 25;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 2:
				stopdraw = 1;
				qlock(rend.l);
				accuracy -= 25;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 3:
				stopdraw = 1;
				qlock(rend.l);
				jcx += 0.01;
				jcy += 0.01;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 4:
				stopdraw = 1;
				qlock(rend.l);
				jcx -= 0.01;
				jcy -= 0.01;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 5:
				stopdraw = 1;
				qlock(rend.l);
				xmin = -2;
				xmax = 1;
				ymin = -1;
				ymax = 1;
				rwakeup(&rend);
				qunlock(rend.l);
				break;
			case 6:
				threadexitsall(nil);
			}
		}
		if(nbrecv(kctl->c, &r))
        		handlekey(r);
	}
}

Cursor reading = {
	{-1, -1},
	{0xff, 0x80, 0xff, 0x80, 0xff, 0x00, 0xfe, 0x00, 
	 0xff, 0x00, 0xff, 0x80, 0xff, 0xc0, 0xef, 0xe0, 
	 0xc7, 0xf0, 0x03, 0xf0, 0x01, 0xe0, 0x00, 0xc0, 
	 0x03, 0xff, 0x03, 0xff, 0x03, 0xff, 0x03, 0xff, },
	{0x00, 0x00, 0x7f, 0x00, 0x7e, 0x00, 0x7c, 0x00, 
	 0x7e, 0x00, 0x7f, 0x00, 0x6f, 0x80, 0x47, 0xc0, 
	 0x03, 0xe0, 0x01, 0xf0, 0x00, 0xe0, 0x00, 0x40, 
	 0x00, 0x00, 0x01, 0xb6, 0x01, 0xb6, 0x00, 0x00, }
};