/*
 * Copyright 2016, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT license.
 *
 * Authors:
 *		Augustin Cavalier <waddlesplash>
 *		kerwizzy
 */
#include "FractalEngine.h"

#include <algorithm>
#include <math.h>

#include <Bitmap.h>
#include <String.h>

#include "Colorsets.h"


//#define TRACE_MANDELBROT_ENGINE
#ifdef TRACE_MANDELBROT_ENGINE
#	include <stdio.h>
#	define TRACE(x...) printf(x)
#else
#	define TRACE(x...)
#endif


FractalEngine::FractalEngine(BHandler* parent, BLooper* looper)
	:
	BLooper("FractalEngine"),
	fWidth(0), fHeight(0),
	fRenderBuffer(NULL),
	fRenderBufferLen(0),
	fSubsampling(2),
	fMessenger(parent, looper),
	fRenderStopping(false),
	fRenderStopped(true),
	fResizing(false),
	fIterations(1024),
	fColorset(Colorset_Royal)
{
	fDoSet = &FractalEngine::DoSet_Mandelbrot;

	fRenderSem = create_sem(0, "RenderSem");
	fRenderStoppedSem = create_sem(0, "RenderStopped");

	system_info info;
	get_system_info(&info);
	fThreadCount = info.cpu_count;

	if (fThreadCount > MAX_RENDER_THREADS)
		fThreadCount = MAX_RENDER_THREADS;

	for (uint8 i = 0; i < fThreadCount; i++) {
		fRenderThreads[i] = spawn_thread(&FractalEngine::RenderThread,
			BString().SetToFormat("RenderThread%d", i).String(),
			B_NORMAL_PRIORITY, this);
		resume_thread(fRenderThreads[i]);
	}
}


FractalEngine::~FractalEngine()
{
	delete_sem(fRenderSem);
	delete_sem(fRenderStoppedSem);
}


void FractalEngine::MessageReceived(BMessage* msg)
{
	switch (msg->what) {
	case MSG_CHANGE_SET:
		switch (msg->GetUInt8("set", 0)) {
		case 0: fDoSet = &FractalEngine::DoSet_Mandelbrot; break;
		case 1: fDoSet = &FractalEngine::DoSet_BurningShip; break;
		case 2: fDoSet = &FractalEngine::DoSet_Tricorn; break;
		case 3: fDoSet = &FractalEngine::DoSet_Julia; break;
		case 4: fDoSet = &FractalEngine::DoSet_OrbitTrap; break;
		case 5: fDoSet = &FractalEngine::DoSet_Multibrot; break;
		}
		break;
	case MSG_SET_PALETTE:
		switch (msg->GetUInt8("palette", 0)) {
		case 0: fColorset = Colorset_Royal; break;
		case 1: fColorset = Colorset_DeepFrost; break;
		case 2: fColorset = Colorset_Frost; break;
		case 3: fColorset = Colorset_Fire; break;
		case 4: fColorset = Colorset_Midnight; break;
		case 5: fColorset = Colorset_Grassland; break;
		case 6: fColorset = Colorset_Lightning; break;
		case 7: fColorset = Colorset_Spring; break;
		case 8: fColorset = Colorset_HighContrast; break;
		}
		break;
	case MSG_SET_ITERATIONS:
		fIterations = msg->GetUInt16("iterations", 0);
		break;
	case MSG_SET_SUBSAMPLING:
		fSubsampling = msg->GetUInt8("subsampling", 1);
		break;

	case MSG_RESIZE: {
		TRACE("Got MSG_RESIZE\n");
		if (fResizing) {
			// Will be true throughout this whole handler. Set false at the end
			TRACE("Breaking out of MSG_RESIZE handler\n");
			break;
		}
		fResizing = true;
		StopRender();

		delete[] fRenderBuffer;
		fRenderBuffer = NULL;

		fWidth = msg->GetUInt16("width", 320);
		fHeight = msg->GetUInt16("height", 240);

		fRenderBufferLen = fWidth * fHeight * 3;
		fRenderBuffer = new uint8[fRenderBufferLen];
		TRACE("New buffer width %u height %u ptr = %p\n",
			  fWidth, fHeight, fRenderBuffer);

		memset(fRenderBuffer, 0, fRenderBufferLen);

		BMessage message(MSG_BUFFER_CREATED);
		fMessenger.SendMessage(&message);

		fResizing = false;
		break;
	}
	case MSG_RENDER: {
		TRACE("Got MSG_RENDER.\n");
		if (fResizing)
			break;

		// Stop the render if one is already running
		StopRender();
		Render(msg->GetDouble("locationX", 0), msg->GetDouble("locationY", 0),
			msg->GetDouble("size", 0.005));
		break;
	}
	case MSG_THREAD_RENDER_COMPLETE: {
		TRACE("Got MSG_THREAD_RENDER_COMPLETE for thread %d\n",
			msg->GetUInt8("thread", 0));

		int32 threadsStopped;
		get_sem_count(fRenderStoppedSem, &threadsStopped);
		TRACE("threadsStopped = %d\n",threadsStopped);
		if (threadsStopped == fThreadCount) {
			TRACE("Done rendering!\n");
			BMessage message(MSG_RENDER_COMPLETE);
			fMessenger.SendMessage(&message);
		}
		break;
	}

	default:
		BLooper::MessageReceived(msg);
		break;
	}
}


void FractalEngine::WriteToBitmap(BBitmap* bitmap)
{
	Lock();
	BSize size = bitmap->Bounds().Size();
	if (size.IntegerWidth() != fWidth || size.IntegerHeight() != fHeight) {
		// some resize happened and now this won't work.
		Unlock();
		return;
	}
	TRACE("Drawing from = %p\n",fRenderBuffer);
	bitmap->ImportBits(fRenderBuffer,
		fRenderBufferLen, fWidth * 3, 0,
		B_RGB24);
	Unlock();
}


void FractalEngine::StopRender()
{
	if (fRenderStopped || fRenderStopping) {
		// if fRenderStopped is true, then render is already stopped,
		// so we can't stop it again!
		// if fStopRender is true, then the fRenderStoppedSem are already
		// trying to be acquired, so stuff would break if we tried to acquire
		// them again.
		return;
	}
	fRenderStopping = true;
	TRACE("Stopping render...\n");
	for (uint i = 0; i < fThreadCount; i++) {
		TRACE("Stopping thread %d\n",i);
		acquire_sem(fRenderStoppedSem);
			// wait till all the threads are stopped...
	}
	int32 threadsStopped;
	get_sem_count(fRenderStoppedSem, &threadsStopped);
	TRACE("stopped sem count after stop = %d\n",threadsStopped);
	fRenderStopping = false;
	fRenderStopped = true;

	TRACE("Render stopped.\n");
}


void FractalEngine::Render(double locationX, double locationY, double size)
{
	if (fRenderStopping)
		debugger("Error: Render shouldn't be called while fRenderStopping = true\n");
	if (!fRenderStopped)
		debugger("Error: Render already running\n");

	fRenderStopped = false;
		// This means that future Render calls will need to call stop render

	fLocationX = locationX;
	fLocationY = locationY;
	fSize = size;

	TRACE("Location (%%.100g): %.100g;%.100g;%.100g\n", fLocationX, fLocationY, fSize);

	for (uint8 i = 0; i < fThreadCount; i++) {
		release_sem(fRenderSem);
	}
}

status_t FractalEngine::RenderThread(void* data)
{
	FractalEngine* engine = static_cast<FractalEngine*>(data);
	thread_id self = find_thread(NULL);
	uint8 threadNum = 0;
	for (uint8 i = 0; i < engine->fThreadCount; i++) {
		if (engine->fRenderThreads[i] == self) {
			threadNum = i;
			break;
		}
	}

	while (true) {
		TRACE("Thread %d awaiting semaphore...\n", threadNum);
		acquire_sem(engine->fRenderSem);
		TRACE("Thread %d got semaphore!\n", threadNum);

		uint16 width = engine->fWidth;
		uint16 height = engine->fHeight;
		uint16 halfHeight = height / 2;
		uint16 threadWidth = width / engine->fThreadCount;

		uint16 startX = threadWidth * threadNum;
		uint16 endX = threadWidth * (threadNum + 1);

		for (uint16 y = 0; y<halfHeight; y++) {
			for (uint16 x = startX; x<endX; x++) {
				engine->RenderPixel(x, halfHeight + y);
					// max y to RenderPixel will be
					// floor(height/2)*2 = height-height%2.
					// Thus there will be a line of blank pixels if height
					// is not even. Is this a problem?
				engine->RenderPixel(x, halfHeight - y - 1);
					// min y to RenderPixel will be
					// halfHeight-(halfHeight-1)-1 = 0
			}

			if (engine->fRenderStopping) {
				TRACE("Thread %d stopping\n", threadNum);

				// Restart the loop to release fRenderStoppedSem and tell
				// the main thread that this thread is stopped, as well as
				// to update width, height, etc.
				break;
			}
		}

		if (!engine->fRenderStopping) {
			// if we got here, then this thread has finished rendering.
			BMessage message(FractalEngine::MSG_THREAD_RENDER_COMPLETE);
			message.AddUInt8("thread", threadNum);
			engine->PostMessage(&message);
		}

		release_sem(engine->fRenderStoppedSem);
	}
	return B_OK;
}


void FractalEngine::RenderPixel(uint32 x, uint32 y)
{
	double real = (x * fSize + fLocationX) - (fWidth / 2 * fSize);
	double imaginary = (y * -(fSize) + fLocationY) - (fHeight / 2 * -(fSize));

	double subsampleDelta = fSize / fSubsampling;
	uint16 nSamples = fSubsampling * fSubsampling;

	int16 totalR = 0;
	int16 totalG = 0;
	int16 totalB = 0;
	for (uint8 subX = 0; subX < fSubsampling; subX++) {
		for (uint8 subY = 0; subY < fSubsampling; subY++) {
			double sampleReal = real + subX * subsampleDelta;
			double sampleImaginary = imaginary + subY * subsampleDelta;

			int32 sampleIterToEscape = (this->*fDoSet)(sampleReal, sampleImaginary);

			uint16 sampleLoc = 0;
			if (sampleIterToEscape == -1)
				// Didn't escape.
				sampleLoc = 999;
			else
				sampleLoc = 998 - (sampleIterToEscape % 999);
			sampleLoc *= 3;

			totalR += fColorset[sampleLoc + 0];
			totalG += fColorset[sampleLoc + 1];
			totalB += fColorset[sampleLoc + 2];
		}
	}


	uint32 offsetBase = fWidth * y * 3 + x * 3;
	fRenderBuffer[offsetBase + 2] = totalR / nSamples; // fRenderBuffer is BGR
	fRenderBuffer[offsetBase + 1] = totalG / nSamples;
	fRenderBuffer[offsetBase + 0] = totalB / nSamples;
}


// Magic numbers & other general constants
const double gJuliaA = 0;
const double gJuliaB = 1;

const uint8 gEscapeHorizon = 4;


int32 FractalEngine::DoSet_Mandelbrot(double real, double imaginary)
{
	double zReal = 0;
	double zImaginary = 0;

	for (int32 i = 0; i < fIterations; i++) {
		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq + (-1 * zImaginarySq));

		zImaginary = 2 * (zReal * zImaginary);
		zReal = nzReal;

		zReal += real;
		zImaginary += imaginary;

		// If it is outside the 2 unit circle...
		if ((zRealSq + zImaginarySq) > gEscapeHorizon)
			return i; // stop it from running longer
	}
	return -1;
}


int32 FractalEngine::DoSet_BurningShip(double real, double imaginary)
{
	double zReal = 0;
	double zImaginary = 0;

	// It looks "upside down" otherwise.
	imaginary = -imaginary;

	for (int32 i = 0; i < fIterations; i++) {
		zReal = fabs(zReal);
		zImaginary = fabs(zImaginary);

		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq + (-1 * zImaginarySq));

		zImaginary = 2 * (zReal * zImaginary);
		zReal = nzReal;

		zReal += real;
		zImaginary += imaginary;

		// If it is outside the 2 unit circle...
		if ((zRealSq + zImaginarySq) > gEscapeHorizon)
			return i; // stop it from running longer
	}
	return -1;
}


int32 FractalEngine::DoSet_Tricorn(double real, double imaginary)
{
	double zReal = 0;
	double zImaginary = 0;

	real = -real;

	for (int32 i = 0; i < fIterations; i++) {
		double znRe = zImaginary * -1;
		zImaginary = zReal * -1;
		zReal = znRe; // Swap the real and complex parts each time.

		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq + (-1 * zImaginarySq));

		zImaginary = 2 * (zReal * zImaginary);
		zReal = nzReal;

		zReal += real;
		zImaginary += imaginary;

		// If it is outside the 2 unit circle...
		if ((zRealSq + zImaginarySq) > gEscapeHorizon)
			return i; // stop it from running longer
	}
	return -1;
}


int32 FractalEngine::DoSet_Julia(double real, double imaginary)
{
	double zReal = real;
	double zImaginary = imaginary;

	double muRe = gJuliaA;
	double muIm = gJuliaB;

	for (int32 i = 0; i < fIterations; i++) {
		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq + (-1 * (zImaginarySq)));

		zImaginary = 2 * (zReal * zImaginary);
		zReal = nzReal;

		zReal += muRe;
		zImaginary += muIm;

		// If it is outside the 2 unit circle...
		if ((zRealSq + zImaginarySq) > gEscapeHorizon)
			return i; // stop it from running longer
	}
	return -1;
}


int32 FractalEngine::DoSet_OrbitTrap(double real, double imaginary)
{
	double zReal = 0;
	double zImaginary = 0;

	double closest = 10000000;
	double distance = 0;
	double lineDist = 0;

	for (int32 i = 0; i < fIterations; i++) {
		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq + (-1 * zImaginarySq));

		zImaginary = 2 * (zReal * zImaginary);
		zReal = nzReal;

		zReal += real;
		zImaginary += imaginary;

		distance = sqrt(zRealSq + zImaginarySq);
		lineDist = fabs(zReal + zImaginary);

		// If it is closer than ever before...
		if (lineDist < closest)
			closest = lineDist;

		if (distance > gEscapeHorizon)
			return static_cast<int32>(floor(4 * log(4 / closest)));
	}
	return static_cast<int32>(floor(4 * log(4 / closest)));
}


int32 FractalEngine::DoSet_Multibrot(double real, double imaginary)
{
	double zReal = 0;
	double zImaginary = 0;

	for (int32 i = 0; i < fIterations; i++) {
		double zRealSq = zReal * zReal;
		double zImaginarySq = zImaginary * zImaginary;
		double nzReal = (zRealSq * zReal - 3 * zReal * zImaginarySq);

		zImaginary = 3 * (zRealSq * zImaginary) - (zImaginarySq * zImaginary);

		zReal = nzReal;
		zReal += real;
		zImaginary += imaginary;

		// If it is outside the 2 unit circle...
		if ((zRealSq + zImaginarySq) > gEscapeHorizon)
			return i; // stop it from running longer
	}
	return -1;
}