#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <math.h>
#include <complex.h>
#include <unistd.h>
#include <pthread.h>
#include <assert.h>
#include "lodepng.h"

#ifndef CMPLX
#define CMPLX(r,i) ((r) + (i)*I)
#endif

const int MAXITER=200;
const double EPS=1e-1;

static double complex att_func(double complex x,double exponent){
	return clog((1-cpow(x,exponent+1))/(1-x));

	// return csqrt(M_SQRT2*x);  // sqrt(2) and 0
	// return csqrt(1+x);  // phi
	// return csqrt(2*x-1);  // attracts REALLY SLOWLY to 1+0i
	// return clog(x);  // two basins and two infinities
	// return clog(x*x);  // "241" basins
	// return clog(x*x+x);  // 3 basins but intricate structure
	// return clog(x*x*x*x+x*x*x+x*x+x+1);  // 2 basin fractal?
	// return csqrt(x*x*x+x*x+x+1);
}

struct colour{
	unsigned char r,g,b;
};

static struct colour make_colour(double complex z){
	static const double N=3;
	double norm=cabs(z);
	if(norm>N)norm=N;
	int normclr=(int)(norm/N*255);
	int argclr=(int)((carg(z)+M_PI)/(2*M_PI)*255);
	return (struct colour){normclr,argclr,0};
}

struct config{
	double unitsz;
	int width,height;
	double midx,midy;
	double exponent;
};

struct thread_info{
	pthread_t th;
	int ystart,yend;
	unsigned char *img;
	bool do_basins;
	struct config cfg;
};

static void* thread_entry(void *arg_vp){
	struct thread_info *th_info=arg_vp;
	int ystart=th_info->ystart;
	int yend=th_info->yend;
	unsigned char *img=th_info->img;
	bool do_basins=th_info->do_basins;
	struct config cfg=th_info->cfg;

	int nbasins=0,bsz=0;
	double complex *basins=NULL;
	if(do_basins){
		bsz=16;
		basins=malloc(bsz*sizeof(double complex));
		assert(basins);
	}

	printf("Thread: y in [%d,%d)\n",ystart,yend);

	for(int iy=ystart;iy<yend;iy++){
		for(int ix=0;ix<cfg.width;ix++){
			const double complex start=CMPLX(
					(double)(ix-cfg.width/2)/cfg.unitsz+cfg.midx,
					cfg.midy-(double)(iy-cfg.height/2)/cfg.unitsz);
			double complex z=start;
			for(int i=0;i<MAXITER;i++){
				z=att_func(z,cfg.exponent);
			}
			if(do_basins){
				int bi;
				for(bi=0;bi<nbasins;bi++){
					if(cabs(z-basins[bi])<EPS)break;
				}
				if(bi==nbasins){
					if(nbasins==bsz){
						bsz*=2;
						basins=realloc(basins,bsz*sizeof(double complex));
						assert(basins);
					}
					fprintf(stderr,"New basin at %lf,%lf (from point %lf,%lf)\n",creal(z),cimag(z),creal(start),cimag(start));
					basins[bi]=z;
					nbasins++;
				}
				basins[bi]=(basins[bi]+z)/2;
				
				// struct colour clr=make_colour(bi);
			}

			struct colour clr=make_colour(z);
			img[3*(iy*cfg.width+ix)+0]=clr.r;
			img[3*(iy*cfg.width+ix)+1]=clr.g;
			img[3*(iy*cfg.width+ix)+2]=clr.b;
		}
	}

	pthread_exit(NULL);
}

int main(int argc,char **argv){
	const double minx=-3,miny=-3;
	const double maxx=3,maxy=3;
	const double width=1000;

	struct config cfg={
		.unitsz=width/(maxx-minx),
		.width=width, .height=width*(maxy-miny)/(maxx-minx),
		.midx=(maxx+minx)/2, .midy=(maxy+miny)/2,
		.exponent=8.5
	};

	if(argc>=2){
		cfg.exponent=strtod(argv[1],NULL);
	}
	const char *fname;
	if(argc>=3){
		fname=argv[2];
	} else {
		fname="out.png";
	}

	printf("Config:\n");
	printf("unitsz = %lf\n",cfg.unitsz);
	printf("width = %d\n",cfg.width);
	printf("height = %d\n",cfg.height);
	printf("midx = %lf\n",cfg.midx);
	printf("midy = %lf\n",cfg.midy);
	printf("exponent = %lf\n",cfg.exponent);

	unsigned char img[cfg.width*cfg.height*3];

	int nthreads=sysconf(_SC_NPROCESSORS_ONLN);
	if(nthreads==4)nthreads/=2;
	if(nthreads>=999||nthreads<1)nthreads=1;
	fprintf(stderr,"Using %d thread%s\n",nthreads,nthreads==1?"":"s");


	pthread_attr_t attr;
	if(pthread_attr_init(&attr)!=0){
		perror("pthread_attr_init");
		return 1;
	}

	struct thread_info ths[nthreads];
	for(int i=0;i<nthreads;i++){
		ths[i].ystart=i*cfg.height/nthreads;
		ths[i].yend=(i+1)*cfg.height/nthreads;
		ths[i].img=img;
		ths[i].do_basins=nthreads==1;
		ths[i].cfg=cfg;
		int ret=pthread_create(&ths[i].th,&attr,thread_entry,&ths[i]);
		if(ret!=0){
			perror("pthread_create");
			return 1;
		}
	}

	pthread_attr_destroy(&attr);

	for(int i=0;i<nthreads;i++){
		int ret=pthread_join(ths[i].th,NULL);
		if(ret!=0){
			perror("pthread_join");
			return 1;
		}
	}

	lodepng_encode_file(fname,img,cfg.width,cfg.height,LCT_RGB,8);
}