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

using namespace psemek;

template <int D>
util::array<std::uint8_t, D> generate(std::array<std::size_t, D> const & size, int octaves, int minoctave, float power,
	bool tile, float gamma, bool remap, std::uint64_t seed)
{
	std::cout << "Generating " << D << "D 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 << "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::array<geom::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::array<float, D> result_float(size);
	geom::interval<float> value_range;

	for (auto idx : result_float.indices())
	{
		geom::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::array<std::uint8_t, D> result(size);

	for (auto idx : result.indices())
	{
		float v = result_float(idx);
		if (remap)
			v = geom::unlerp(value_range, 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:

    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

    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::uint64_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[])
{
	std::unordered_map<std::string, std::string> values;
	values["dim"] = "2";
	values["sizex"] = "256";
	values["sizey"] = "256";
	values["sizez"] = "256";
	values["octaves"] = "8";
	values["minoctave"] = "0";
	values["power"] = "1.0";
	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 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 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 (!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 (!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)
	{
		auto result = generate<2>({*sizex, *sizey}, *octaves, *minoctave, *power, *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)
	{
		auto result = generate<3>({*sizex, *sizey, *sizez}, *octaves, *minoctave, *power,  *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());
	}
}