#include <iostream>
#include <fstream>
#include <vector>
#include <array>
#include <thread>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <cstdlib>
#include <cmath>
#include <mpfr.h>
#include "lodepng/lodepng.h"


#define MULTITHREAD


constexpr size_t precision = 128;
constexpr mpfr_rnd_t rounding = MPFR_RNDN;


struct colorscheme_params {
	double speed = 1.0;
	double offset = 0.0;
};

static std::array<uint8_t, 3> colorscheme(double value, const colorscheme_params &params) {
	if (value == -1) return {0, 0, 0};
	else {
		// 720 is a full circle
		const double hue = params.offset + params.speed * value / 2;

		// https://en.wikipedia.org/wiki/HSL_and_HSV#HSL_to_RGB_alternative
		// saturation = 1.0
		// lightness = 0.5
		// a = saturation * min(lightness, 1 - lightness) = 0.5
		const auto f = [hue](int n) {
			// k = (n + hue / 30°) mod 12
			// f(n) = lightness - a * max(-1, min(k - 3, 9 - k, 1))
			double k = std::fmod(n + hue / 30, 12);
			double res = 0.5 - 0.5 * std::max(-1.0, std::min(1.0, std::min(k - 3, 9 - k)));

			return (uint8_t)(res * 127.99 + 127.99);
		};

		return {f(0), f(8), f(4)};
	}
}

struct window {
	double xmin, xmax, ymin, ymax;

	window enlarge_to_ratio(double width_over_height) const {
		window res = *this;

		double current_rat = (xmax - xmin) / (ymax - ymin);
		if (current_rat > width_over_height) {
			double h = (xmax - xmin) / width_over_height;
			double ymid = (ymin + ymax) / 2;
			res.ymin = ymid - h / 2;
			res.ymax = ymid + h / 2;
		} else if (current_rat < width_over_height) {
			double w = (ymax - ymin) * width_over_height;
			double xmid = (xmin + xmax) / 2;
			res.xmin = xmid - w / 2;
			res.xmax = xmid + w / 2;
		}

		return res;
	}

	window recenter_x(double x) const {
		window res = *this;
		res.xmin = x - (xmax - xmin) / 2;
		res.xmax = x + (xmax - xmin) / 2;
		return res;
	}

	window recenter_y(double y) const {
		window res = *this;
		res.ymin = y - (ymax - ymin) / 2;
		res.ymax = y + (ymax - ymin) / 2;
		return res;
	}

	window set_width(double w) const {
		window res = *this;
		res.xmin = (xmax + xmin) / 2 - w / 2;
		res.xmax = (xmax + xmin) / 2 + w / 2;
		return res;
	}

	window set_height(double h) const {
		window res = *this;
		res.ymin = (ymax + ymin) / 2 - h / 2;
		res.ymax = (ymax + ymin) / 2 + h / 2;
		return res;
	}
};

class mpfr_class {
public:
	mpfr_class() {
		mpfr_init2(num, precision);
	}

	~mpfr_class() {
		mpfr_clear(num);
	}

	void set(const mpfr_class &a) {
		mpfr_set(num, a.num, rounding);
	}

	void set_d(double a) {
		mpfr_set_d(num, a, rounding);
	}

	void add(const mpfr_class &a, const mpfr_class &b) {
		mpfr_add(num, a.num, b.num, rounding);
	}

	void sub(const mpfr_class &a, const mpfr_class &b) {
		mpfr_sub(num, a.num, b.num, rounding);
	}

	void mul(const mpfr_class &a, const mpfr_class &b) {
		mpfr_mul(num, a.num, b.num, rounding);
	}

	void mul_d(const mpfr_class &a, double b) {
		mpfr_mul_d(num, a.num, b, rounding);
	}

	void fma(const mpfr_class &a, const mpfr_class &b, const mpfr_class &c) {
		mpfr_fma(num, a.num, b.num, c.num, rounding);
	}

	void sqrt(const mpfr_class &a) {
		mpfr_sqrt(num, a.num, rounding);
	}

	void log(const mpfr_class &a) {
		mpfr_log(num, a.num, rounding);
	}

	void log2(const mpfr_class &a) {
		mpfr_log2(num, a.num, rounding);
	}

	int cmp_ui(unsigned long int a) const {
		return mpfr_cmp_ui(num, a);
	}

	bool greater(const mpfr_class &a) const {
		return mpfr_greater_p(num, a.num);
	}

	double get_d() const {
		return mpfr_get_d(num, rounding);
	}

private:
	mpfr_t num;
};

template <typename T>
class mpfr_float_shim {
public:
	void set(const mpfr_float_shim &a) {
		num = a.num;
	}

	void set_d(double a) {
		num = a;
	}

	void add(const mpfr_float_shim &a, const mpfr_float_shim &b) {
		num = a.num + b.num;
	}

	void sub(const mpfr_float_shim &a, const mpfr_float_shim &b) {
		num = a.num - b.num;
	}

	void mul(const mpfr_float_shim &a, const mpfr_float_shim &b) {
		num = a.num * b.num;
	}

	void mul_d(const mpfr_float_shim &a, double b) {
		num = a.num * b;
	}

	void fma(const mpfr_float_shim &a, const mpfr_float_shim &b, const mpfr_float_shim &c) {
		num = a.num * b.num + c.num;
	}

	void sqrt(const mpfr_float_shim &a) {
		num = std::sqrt(a.num);
	}

	void log(const mpfr_float_shim &a) {
		num = std::log(a.num);
	}

	void log2(const mpfr_float_shim &a) {
		num = std::log2(a.num);
	}

	bool greater(const mpfr_float_shim &a) const {
		return num > a.num;
	}

	int cmp_ui(unsigned long int a) const {
		if (num < a) return -1;
		if (num > a) return 1;
		return 0;
	}

	double get_d() const {
		return num;
	}

private:
	T num;
};

using FLT = mpfr_float_shim<double>;

// Destroys the values in both escape_x and escape_y
static double mandel_smooth_iters(size_t iters, FLT escape_x, FLT escape_y, int log_log_bound) {
	escape_y.mul(escape_y, escape_y);
	escape_x.fma(escape_x, escape_x, escape_y);
	escape_x.log2(escape_x);
	escape_x.log2(escape_x);
	return iters + 1 + log_log_bound - escape_x.get_d();
}

struct params {
	std::vector<float> *image;
	size_t width, height;
	size_t maxiter;
	const FLT &xmin, &xmax, &ymin, &ymax;
};

static void mandel_job(const params &params, size_t from_idx, size_t to_idx) {
	FLT x, y, a, b, a2, b2, magnitude;

	// bound for smooth iteration count
	// (source: https://www.iquilezles.org/www/articles/mset_smooth/mset_smooth.htm)
	FLT bound_squared;
	int log_log_bound = 3;
	bound_squared.set_d(1 << (2 << log_log_bound));  // (2^(2^k))^2 = 2^(2^k + 2^k) = 2^(2 * 2^k)

	for (size_t index = from_idx; index < to_idx; index++) {
		const size_t xi = index % params.width;
		const size_t yi = index / params.width;

		x.sub(params.xmax, params.xmin);
		x.mul_d(x, (double)xi / (params.width - 1));
		x.add(x, params.xmin);

		// min and max reversed here because the vertical axis is mirrored
		y.sub(params.ymin, params.ymax);
		y.mul_d(y, (double)yi / (params.height - 1));
		y.add(y, params.ymax);

		a.set(x);
		b.set(y);
		a2.mul(a, a);
		b2.mul(b, b);

		size_t iter;
		for (iter = 0; iter < params.maxiter; iter++) {
			b.mul(a, b);
			b.add(b, b);
			b.add(b, y);

			a.sub(a2, b2);
			a.add(a, x);

			a2.mul(a, a);
			b2.mul(b, b);

			magnitude.add(a2, b2);

			if (magnitude.greater(bound_squared)) break;
		}

		const double smooth_iters =
			iter == params.maxiter
				? -1
				: mandel_smooth_iters(iter, a, b, log_log_bound);

		(*params.image)[index] = smooth_iters;
	}
}

static std::vector<uint8_t> colorise_image(const std::vector<float> &image, const colorscheme_params &clrpar) {
	std::vector<uint8_t> pixels(3 * image.size());

	for (size_t i = 0; i < image.size(); i++) {
		const double iters = image[i];

		std::array<uint8_t, 3> clrs = colorscheme(iters, clrpar);
		pixels[3 * i + 0] = clrs[0];
		pixels[3 * i + 1] = clrs[1];
		pixels[3 * i + 2] = clrs[2];
	}

	return pixels;
}

struct job {
	size_t from_idx, to_idx;
};

struct job_list {
	std::vector<job> jobs;
	std::atomic_size_t jobidx;
};

static void mandel_worker(const params &params, job_list &jl) {
	while (true) {
		size_t ji = jl.jobidx.fetch_add(1);
		if (ji >= jl.jobs.size()) break;
		mandel_job(params, jl.jobs[ji].from_idx, jl.jobs[ji].to_idx);
	}
}

static void write_pixels_to_file(const std::string &output_fname, const std::vector<uint8_t> &pixels, size_t width, size_t height) {
	int ret = lodepng::encode(output_fname, pixels, width, height, LCT_RGB, 8);
	if (ret != 0) {
		std::cerr << "Cannot write PNG: " << lodepng_error_text(ret) << std::endl;
		exit(1);
	}
}

int main(int argc, char **argv) {
	size_t width = 1920, height = 1080;
	size_t maxiter = 256;
	window win{-2.9, 1.9, -1.35, 1.35};
	std::string save_fname, load_fname, output_fname;
	colorscheme_params clrpar;

	for (int i = 1; i < argc; i++) {
		const auto take_argument = [&i, argc, argv](auto &&parsefunc, auto &dest) {
			if (i == argc - 1) {
				std::cerr << "Flag expects argument: '" << argv[i] << "'" << std::endl;
				exit(1);
			}

			dest = parsefunc(argv[i + 1]);
			i++;
		};

		const auto &parse_int = [](const char *s) { return std::stoull(s); };
		const auto &parse_flt = [](const char *s) { return std::stod(s); };
		const auto &identity = [](const char *s) { return s; };

		if      (strcmp(argv[i], "-W") == 0) take_argument(parse_int, width);
		else if (strcmp(argv[i], "-H") == 0) take_argument(parse_int, height);
		else if (strcmp(argv[i], "-M") == 0) take_argument(parse_int, maxiter);
		else if (strcmp(argv[i], "-o") == 0) take_argument(identity, output_fname);
		else if (strcmp(argv[i], "-s") == 0) take_argument(identity, save_fname);
		else if (strcmp(argv[i], "-l") == 0) take_argument(identity, load_fname);
		else if (strcmp(argv[i], "-L") == 0) take_argument(parse_flt, win.xmin);
		else if (strcmp(argv[i], "-R") == 0) take_argument(parse_flt, win.xmax);
		else if (strcmp(argv[i], "-T") == 0) take_argument(parse_flt, win.ymin);
		else if (strcmp(argv[i], "-B") == 0) take_argument(parse_flt, win.ymax);
		else if (strcmp(argv[i], "-X") == 0) {
			double x; take_argument(parse_flt, x);
			win = win.recenter_x(x);
		}
		else if (strcmp(argv[i], "-Y") == 0) {
			double y; take_argument(parse_flt, y);
			win = win.recenter_y(y);
		}
		else if (strcmp(argv[i], "-Z") == 0) {
			double w; take_argument(parse_flt, w);
			win = win.set_width(w).set_height((double)height / width * w);
		}
		else if (strcmp(argv[i], "-cS") == 0) take_argument(parse_flt, clrpar.speed);
		else if (strcmp(argv[i], "-cO") == 0) take_argument(parse_flt, clrpar.offset);
		else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
			std::cout <<
				"Usage: " << argv[0] << " [options] -o <out.png>\n"
				"       " << argv[0] << " [options] -s <data.txt>\n"
				"       " << argv[0] << " [options] -l <data.txt> -o <out.png>\n"
				"\n"
				"  -h/--help      Show this help\n"
				"\n"
				"  -W <width>     Width of output image in pixels [1920]\n"
				"  -H <height>    Height of output image in pixels [1080]\n"
				"  -M <maxiter>   Number of iterations after which a point is \"in set\" [256]\n"
				"  -o <out.png>   Output file name\n"
				"  -s <data.txt>  Iteration cache output file name (colorscheme options unused)\n"
				"  -l <data.txt>  Only colorise an iteration cache file; requires -o (and\n"
				"                 ignores all other flags except colorscheme flags)\n"
				"\n"
				"For the complex boundaries, either set them explicitly (they will potentially\n"
				"be expanded so that the width/height ratio is satisfied):\n"
				"  -L <left>      Left real bound [-2.9]\n"
				"  -R <right>     Right real bound [1.9]\n"
				"  -B <bottom>    Bottom imaginary bound [-1.35]\n"
				"  -T <top>       Top imaginary bound [1.35]\n"
				"or via a center point and zoom level:\n"
				"  -X <midx>      Real part of midpoint of image [-0.5]\n"
				"  -Y <midy>      Imaginary part of midpoint of image [0.0]\n"
				"  -Z <cplxwid>   Complex width of the image (right - left) [4.8]\n"
				"When mixing these two approaches, later arguments override earlier ones.\n"
				"\n"
				"Colorscheme options: (unused for -s)\n"
				"  -cS <speed>    Hue spin speed [1.0]\n"
				"  -cO <offset>   Hue spin offset [0.0]\n";
			std::cout << std::flush;
			return 0;
		} else {
			std::cerr << "Unrecognised option '" << argv[i] << "'" << std::endl;
			return 1;
		}
	}

	if ((!save_fname.empty() && (!load_fname.empty() || !output_fname.empty())) ||
			(save_fname.empty() && output_fname.empty())) {
		std::cerr << "Either -o, -s or -l -o is required. Use -h for help." << std::endl;
		return 1;
	}

	if (!load_fname.empty()) {
		std::ifstream f(load_fname, std::ios::binary);
		if (!f) {
			std::cerr << "Cannot open file '" << load_fname << "'" << std::endl;
			return 1;
		}
		f.read((char*)&width, sizeof width);
		f.read((char*)&height, sizeof height);

		std::vector<float> image(width * height);
		for (size_t i = 0; i < width * height; i++) {
			f.read((char*)&image[i], sizeof(float));
		}

		std::vector<uint8_t> pixels = colorise_image(image, clrpar);
		write_pixels_to_file(output_fname, pixels, width, height);
		return 0;
	}

	if (width == 0 || height == 0 || maxiter == 0 ||
			win.xmin > win.xmax || win.ymin > win.ymax) {
		std::cerr << "Invalid parameters given" << std::endl;
		return 1;
	}

	win = win.enlarge_to_ratio((double)width / height);

	FLT xmin, xmax, ymin, ymax;
	xmin.set_d(win.xmin);
	xmax.set_d(win.xmax);
	ymin.set_d(win.ymin);
	ymax.set_d(win.ymax);

	std::vector<float> image(width * height);

	params params{
		&image,
		width, height,
		maxiter,
		xmin, xmax, ymin, ymax
	};

#ifdef MULTITHREAD
	const size_t nthreads = std::thread::hardware_concurrency();
	const size_t njobs = 16 * nthreads;

	job_list jl;
	jl.jobs.resize(njobs);
	for (size_t i = 0; i < njobs; i++) {
		jl.jobs[i].from_idx = i * width * height / njobs;
		jl.jobs[i].to_idx = (i + 1) * width * height / njobs;
	}
	jl.jobidx.store(0);

	std::vector<std::thread> threads;
	for (size_t i = 0; i < nthreads; i++) {
		threads.emplace_back([&params, &jl, i]() {
			mandel_worker(params, jl);
		});
	}

	for (std::thread &th : threads) th.join();
#else
	mandel_job(params, 0, width * height);
#endif

	if (!save_fname.empty()) {
		std::ofstream f(save_fname, std::ios::binary);
		if (!f) {
			std::cerr << "Cannot open file '" << save_fname << "'" << std::endl;
			return 1;
		}

		f.write((const char*)&width, sizeof width);
		f.write((const char*)&height, sizeof height);
		for (size_t i = 0; i < width * height; i++) {
			f.write((const char*)&image[i], sizeof(float));
		}
	} else {
		std::vector<uint8_t> pixels = colorise_image(image, clrpar);
		write_pixels_to_file(output_fname, pixels, width, height);
	}
}