#include <psemek/pcg/perlin.hpp>
#include <psemek/pcg/fractal.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/random/uniform_sphere.hpp>
#include <psemek/gfx/pixmap.hpp>
#include <psemek/util/hash_table.hpp>
#include <psemek/io/file_stream.hpp>

using namespace psemek;

template <int D>
auto moore_neighbourhood()
{
	if constexpr (D == 2)
	{
		std::array<math::vector<int, D>, 9> result;
		for (int dy = -1; dy <= 1; ++dy)
			for (int dx = -1; dx <= 1; ++dx)
				result[(dy + 1) * 3 + (dx + 1)] = {dx, dy};
		return result;
	}
	else if constexpr (D == 3)
	{
		std::array<math::vector<int, D>, 27> result;
		for (int dz = -1; dz <= 1; ++dz)
			for (int dy = -1; dy <= 1; ++dy)
				for (int dx = -1; dx <= 1; ++dx)
					result[(dz + 1) * 9 + (dy + 1) * 3 + (dx + 1)] = {dx, dy, dz};
		return result;
	}
}

template <int D>
util::ndarray<std::uint8_t, D> generate_perlin(std::array<std::size_t, D> const & size, int octaves, int minoctave, float power,
	bool invert, bool tile, float gamma, bool remap, std::uint64_t seed)
{
	std::cout << "Generating " << D << "D perlin noise\n";
	std::cout << "Size ";
	for (int d = 0; d < D; ++d)
	{
		if (d > 0)
			std::cout << "x";
		std::cout << size[d];
	}

	std::cout << "\nOctaves: " << minoctave << ".." << (minoctave + octaves - 1) << "\n";
	std::cout << "Power: " << power << "\n";
	std::cout << "Invert: " << (invert ? "true" : "false") << "\n";
	std::cout << "Tile: " << (tile ? "true" : "false") << "\n";
	std::cout << "Gamma: " << gamma << "\n";
	std::cout << "Remap: " << (remap ? "true" : "false") << "\n";
	std::cout << "Seed: " << std::hex << std::setw(16) << std::setfill('0') << seed << "\n";
	std::cout << std::flush;

	random::generator rng{seed, 0};

	std::vector<pcg::perlin<float, D>> octave_noise;
	std::vector<float> octave_weights;
	float sum_octave_weights = 0.f;

	random::uniform_sphere_vector_distribution<float, D> random_gradient;

	for (int o = minoctave; o < (minoctave + octaves); ++o)
	{
		std::array<std::size_t, D> octave_size;
		for (int d = 0; d < D; ++d)
			octave_size[d] = (1 << o) + (tile ? 0 : 1);

		util::ndarray<math::vector<float, D>, D> gradients(octave_size);
		for (auto & g : gradients)
			g = random_gradient(rng);

		if (tile)
			octave_noise.emplace_back(std::move(gradients), pcg::seamless);
		else
			octave_noise.emplace_back(std::move(gradients));

		float weight = std::pow(static_cast<float>(1 << o), - power);
		octave_weights.push_back(weight);
		sum_octave_weights += weight;
	}

	for (auto & w : octave_weights)
		w /= sum_octave_weights;

	pcg::fractal<pcg::perlin<float, D>> fractal_noise(std::move(octave_noise), std::move(octave_weights));

	int largest_size = size[0];
	for (int d = 0; d < D; ++d)
		largest_size = std::max<int>(largest_size, size[d]);

	util::ndarray<float, D> result_float(size);
	math::interval<float> value_range;

	for (auto idx : result_float.indices())
	{
		math::point<float, D> p;

		for (int d = 0; d < D; ++d)
		{
			float unused;
			p[d] = std::modf((idx[d] + 0.5f) / largest_size, &unused);
		}

		float v = fractal_noise(p);

		result_float(idx) = v;
		value_range |= v;
	}

	util::ndarray<std::uint8_t, D> result(size);

	for (auto idx : result.indices())
	{
		float v = result_float(idx);
		if (remap)
			v = math::unlerp(value_range, v);
		if (invert)
			v = 1.f - v;
		v = std::pow(v, gamma);
		result(idx) = static_cast<std::uint8_t>(std::clamp(v * 255.f, 0.f, 255.f));
	}

	return result;
}

template <int D>
util::ndarray<std::uint8_t, D> generate_worley(std::array<std::size_t, D> const & size, int octaves, int minoctave, float power,
	bool invert, bool tile, float gamma, bool remap, std::uint64_t seed)
{
	std::cout << "Generating " << D << "D worley noise\n";
	std::cout << "Size ";
	for (int d = 0; d < D; ++d)
	{
		if (d > 0)
			std::cout << "x";
		std::cout << size[d];
	}

	std::cout << "\nOctaves: " << minoctave << ".." << (minoctave + octaves - 1) << "\n";
	std::cout << "Power: " << power << "\n";
	std::cout << "Invert: " << (invert ? "true" : "false") << "\n";
	std::cout << "Tile: " << (tile ? "true" : "false") << "\n";
	std::cout << "Gamma: " << gamma << "\n";
	std::cout << "Remap: " << (remap ? "true" : "false") << "\n";
	std::cout << "Seed: " << std::hex << std::setw(16) << std::setfill('0') << seed << "\n";
	std::cout << std::flush;

	random::generator rng{seed, 0};

	std::vector<util::ndarray<math::point<float, D>, D>> octave_points;
	std::vector<float> octave_weights;
	float sum_octave_weights = 0.f;

	for (int o = minoctave; o < minoctave + octaves; ++o)
	{
		std::array<std::size_t, D> octave_size;
		for (int d = 0; d < D; ++d)
			octave_size[d] = (1 << o);

		util::ndarray<math::point<float, D>, D> octave(octave_size);

		for (auto idx : octave.indices())
		{
			auto & p = octave(idx);
			for (int d = 0; d < D; ++d)
				p[d] = idx[d] + random::uniform<float>(rng);
		}

		octave_points.push_back(std::move(octave));

		float weight = std::pow(static_cast<float>(1 << o), - power);
		octave_weights.push_back(weight);
		sum_octave_weights += weight;
	}

	for (auto & w : octave_weights)
		w /= sum_octave_weights;

	util::ndarray<float, D> result_float(size);
	math::interval<float> value_range;

	for (auto idx : result_float.indices())
	{
		math::point<float, D> p;

		for (int d = 0; d < D; ++d)
			p[d] = (idx[d] + 0.5f) / size[d] * (1 << minoctave);

		float v = 0.f;

		for (int o = 0; o < octaves; ++o)
		{
			math::point<int, D> i;
			for (int d = 0; d < D; ++d)
				i[d] = std::floor(p[d]);

			int octave_size = 1 << (o + minoctave);

			float closest_distance = std::numeric_limits<float>::infinity();

			for (auto n : moore_neighbourhood<D>())
			{
				math::point<int, D> j = i + n;

				if (tile)
				{
					for (int d = 0; d < D; ++d)
						j[d] = (j[d] + octave_size) % octave_size;
				}
				else
				{
					bool good = true;
					for (int d = 0; d < D; ++d)
					{
						good &= (j[d] >= 0);
						good &= (j[d] < octave_size);
					}

					if (!good)
						continue;
				}

				math::vector<float, D> r = octave_points[o](j) - p;

				if (tile)
				{
					for (int d = 0; d < D; ++d)
					{
						if (i[d] + n[d] == -1)
							r[d] -= octave_size;
						else if (i[d] + n[d] == octave_size)
							r[d] += octave_size;
					}
				}

				math::make_min(closest_distance, math::length(r));
			}

			v += closest_distance * octave_weights[o] / std::sqrt(2.f);

			for (int d = 0; d < D; ++d)
				p[d] *= 2.f;
		}

		result_float(idx) = v;
		value_range |= v;
	}

	util::ndarray<std::uint8_t, D> result(size);

	for (auto idx : result.indices())
	{
		float v = result_float(idx);
		if (remap)
			v = math::unlerp(value_range, v);
		if (invert)
			v = 1.f - v;
		v = std::pow(v, gamma);
		result(idx) = static_cast<std::uint8_t>(std::clamp(v * 255.f, 0.f, 255.f));
	}

	return result;
}

char const options_usage[] =
R"(
Available options:

    type
            Noise type (perlin or worley)
            Default = worley

    dim
            Noise dimension (2 or 3)
            Default = 2

    sizex
            Size along X coordinate, in pixels
            Default = 256

    sizey
            Size along Y coordinate, in pixels
            Default = 256

    sizez
            Size along Z coordinate, in pixels (for dim = 3)
            Default = 256

    octaves
            Number of noise octaves (higher octaves have higher frequency)
            Default = 8

    minoctave
            First octave index (0-th octave is square with the size of largest output dimension)
            Default = 0

    power
            Octaves weight power (0.0 means all octaves have equal contributions), can be negative
            Default = 1.0

    invert
            Replace noise values with 1-value (makes sense for worley noise only)
            Default = false

    tile
            Whether to make the noise tileable (seamless)
            Default = true

    gamma
            The gamma value to apply to the output (gamma = 2.2 means sRGB-like conversion)
            Default = 1.0

    remap
            Whether to remap noise values to [0, 1] range (adding many actaves tends to average out the noise)
            Default = true

    seed
            Random seed (a 64-bit decimal or hex integer, or "random")
            Default = random

)";

std::optional<int> parse_int(std::string const & str)
{
	try
	{
		return std::stoi(str);
	}
	catch (...)
	{
		return std::nullopt;
	}
}

std::optional<std::uint64_t> parse_seed(std::string const & str, int base)
{
	try
	{
		std::size_t unused;
		return std::stoull(str, &unused, base);
	}
	catch (...)
	{
		return std::nullopt;
	}
}

std::optional<float> parse_float(std::string const & str)
{
	try
	{
		return std::stof(str);
	}
	catch (...)
	{
		return std::nullopt;
	}
}

std::optional<bool> parse_bool(std::string const & str)
{
	if (str == "true")
		return true;
	else if (str == "false")
		return false;
	else
		return std::nullopt;
}

int main(int argc, char * argv[])
{
	util::hash_map<std::string, std::string> values;
	values["type"] = "perlin";
	values["dim"] = "2";
	values["sizex"] = "256";
	values["sizey"] = "256";
	values["sizez"] = "256";
	values["octaves"] = "8";
	values["minoctave"] = "0";
	values["power"] = "1.0";
	values["invert"] = "false";
	values["tile"] = "true";
	values["gamma"] = "1.0";
	values["remap"] = "true";
	values["seed"] = "random";

	std::string output_path;

	if (argc < 2)
	{
		std::cerr << "Usage: " << argv[0] << " <output-path> [ key1=value1 [ key2=value2 ... ] ]\n";
		std::cerr << options_usage;
		return EXIT_FAILURE;
	}

	output_path = argv[1];

	for (int i = 2; i < argc; ++i)
	{
		std::string arg = argv[i];
		auto divider = arg.find('=');
		if (divider == std::string::npos) {
			std::cerr << "Failed to parse argument \"" << arg << "\"\n";
			return EXIT_FAILURE;
		}

		auto key = arg.substr(0, divider);
		auto value = arg.substr(divider + 1);

		if (!values.contains(key)) {
			std::cerr << "Unknown key \"" << key << "\"\n";
			return EXIT_FAILURE;
		}

		values[key] = value;
	}

	auto type = values["type"];
	auto dim = parse_int(values["dim"]);
	auto sizex = parse_int(values["sizex"]);
	auto sizey = parse_int(values["sizey"]);
	auto sizez = parse_int(values["sizez"]);
	auto octaves = parse_int(values["octaves"]);
	auto minoctave = parse_int(values["minoctave"]);
	auto power = parse_float(values["power"]);
	auto invert = parse_bool(values["invert"]);
	auto tile = parse_bool(values["tile"]);
	auto gamma = parse_float(values["gamma"]);
	auto remap = parse_bool(values["remap"]);

	std::optional<std::size_t> seed;
	if (values["seed"] == "random") {
		random::device device{};
		seed = (static_cast<std::uint64_t>(device()) << 32) | static_cast<std::uint64_t>(device());
	}
	if (!seed) seed = parse_seed(values["seed"], 10);
	if (!seed) seed = parse_seed(values["seed"], 16);

	if (type != "perlin" && type != "worley")
	{
		std::cerr << "Unknown noise type: " << values["type"] << "\n";
		return EXIT_FAILURE;
	}

	if (!dim || (*dim != 2 && *dim != 3)) {
		std::cerr << "Unknown noise dimension: " << values["dim"] << "\n";
		return EXIT_FAILURE;
	}

	if (!sizex || *sizex <= 0) {
		std::cerr << "Size must be a positive integer: " << values["sizex"] << "\n";
		return EXIT_FAILURE;
	}

	if (!sizey || *sizey <= 0) {
		std::cerr << "Size must be a positive integer: " << values["sizey"] << "\n";
		return EXIT_FAILURE;
	}

	if (!sizez || *sizez <= 0) {
		std::cerr << "Size must be a positive integer: " << values["sizez"] << "\n";
		return EXIT_FAILURE;
	}

	if (!octaves || *octaves <= 0) {
		std::cerr << "Number of octaves must be a positive integer: " << values["octaves"] << "\n";
		return EXIT_FAILURE;
	}

	if (!minoctave || *minoctave < 0) {
		std::cerr << "First octave must be a nonnegative integer: " << values["minoctave"] << "\n";
		return EXIT_FAILURE;
	}

	if (!power) {
		std::cerr << "Power must be a floating-point number: " << values["power"] << "\n";
		return EXIT_FAILURE;
	}

	if (!invert) {
		std::cerr << "Invert must be true or false: " << values["invert"] << "\n";
		return EXIT_FAILURE;
	}

	if (!tile) {
		std::cerr << "Tile must be true or false: " << values["tile"] << "\n";
		return EXIT_FAILURE;
	}

	if (!gamma) {
		std::cerr << "Gamma must be a floating-point number: " << values["gamma"] << "\n";
		return EXIT_FAILURE;
	}

	if (!remap) {
		std::cerr << "Remap must be true or false: " << values["remap"] << "\n";
		return EXIT_FAILURE;
	}

	if (!seed) {
		std::cerr << "Seed must be a decimal or hex integer: " << values["seed"] << "\n";
		return EXIT_FAILURE;
	}

	if (*dim == 2)
	{
		util::ndarray<std::uint8_t, 2> result;
		if (type == "perlin")
			result = generate_perlin<2>({*sizex, *sizey}, *octaves, *minoctave, *power, *invert, *tile, *gamma, *remap, *seed);
		else if (type == "worley")
			result = generate_worley<2>({*sizex, *sizey}, *octaves, *minoctave, *power, *invert, *tile, *gamma, *remap, *seed);

		io::file_ostream out{std::filesystem::path(output_path)};

		if (output_path.ends_with(".png"))
		{
			std::cout << "Saving PNG image to " << output_path << std::endl;
			gfx::write_image_png(result, std::move(out));
		}
		else if (output_path.ends_with(".tga"))
		{
			std::cout << "Saving TGA image to " << output_path << std::endl;
			gfx::write_image_tga(result, std::move(out));
		}
		else if (output_path.ends_with(".pgm"))
		{
			std::cout << "Saving PGM image to " << output_path << std::endl;
			gfx::write_image_pgm(result, std::move(out));
		}
		else
		{
			std::cout << "Saving raw image to " << output_path << std::endl;
			out.write(reinterpret_cast<char const *>(result.data()), result.size());
		}
	}

	if (*dim == 3)
	{
		util::ndarray<std::uint8_t, 3> result;
		if (type == "perlin")
			generate_perlin<3>({*sizex, *sizey, *sizez}, *octaves, *minoctave, *power, *invert, *tile, *gamma, *remap, *seed);
		else if (type == "worley")
			result = generate_worley<3>({*sizex, *sizey, *sizez}, *octaves, *minoctave, *power, *invert, *tile, *gamma, *remap, *seed);

		io::file_ostream out{std::filesystem::path(output_path)};
		std::cout << "Saving raw image to " << output_path << std::endl;
		out.write(reinterpret_cast<char const *>(result.data()), result.size());
	}
}