#include <psemek/app/app.hpp>
#include <psemek/app/main.hpp>
#include <psemek/gfx/mesh.hpp>
#include <psemek/gfx/program.hpp>
#include <psemek/geom/camera.hpp>
#include <psemek/geom/rotation.hpp>
#include <psemek/geom/gram_schmidt.hpp>
#include <psemek/gfx/color.hpp>
#include <psemek/geom/math.hpp>
#include <psemek/geom/interval.hpp>
#include <psemek/random/uniform.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/util/recursive.hpp>

#include <vector>

using namespace psemek;

// All spacial parameters in meters
// All angles in degrees
struct tree_species_description
{
	struct curve_description
	{
		geom::interval<float> base_radius;
		geom::interval<float> segment_length;
		std::function<geom::interval<int>(float)> segment_count;
		geom::interval<float> lean_angle;
		geom::interval<float> rotation_angle;
		float up_tendency;

		float splitting_probability;
		geom::interval<float> splitting_angle;

		geom::interval<float> initial_branching_distance;
		geom::interval<float> branching_distance;
		geom::interval<int> branching_children_count;
		geom::interval<float> branching_rotation;
		geom::interval<float> branching_angle;
		geom::interval<float> branching_size_multiplier;
	};

	curve_description trunk;
	curve_description branch;
};

auto fir_species()
{
	tree_species_description desc;

	desc.trunk.base_radius = {0.1f, 0.15f};
	desc.trunk.segment_length = {0.5f, 0.8f};
	desc.trunk.segment_count = [](float){
		return geom::interval{7, 10};
	};
	desc.trunk.lean_angle = {-5.f, 5.f};
	desc.trunk.rotation_angle = {-5.f, 5.f};
	desc.trunk.up_tendency = 1.f;
	desc.trunk.initial_branching_distance = {0.5f, 0.5f};
	desc.trunk.branching_distance = {0.4f, 0.5f};
	desc.trunk.branching_children_count = {4, 4};
	desc.trunk.branching_rotation = {30, 60.f};
	desc.trunk.branching_angle = {55.f, 65.f};
	desc.trunk.branching_size_multiplier = {1.f, 1.f};
	desc.trunk.splitting_probability = 0.f;
	desc.branch.base_radius = {0.01f, 0.05f};
	desc.branch.segment_length = {0.05f, 0.20f};
	desc.branch.segment_count = [](float h){
		int min = std::max<int>(1, 18 * (1.f - h));
		int max = std::min<int>(20, 20 * (1.f - h));

		if (min > max)
			std::swap(min, max);

		return geom::interval<int>{min, max};
	};
	desc.branch.lean_angle = {-5.f, 5.f};
	desc.branch.rotation_angle = {-1.f, 1.f};
	desc.branch.up_tendency = 0.0f;
	desc.branch.initial_branching_distance = {0.4f, 0.4f};
	desc.branch.branching_distance = {0.1f, 0.2f};
	desc.branch.branching_children_count = {1, 1};
	desc.branch.branching_rotation = {175.f, 185.f};
	desc.branch.branching_angle = {45.f, 45.f};
	desc.branch.branching_size_multiplier = {0.8f, 0.8f};
	desc.branch.splitting_probability = 0.f;

	return desc;
}

auto oak_species()
{
	tree_species_description desc;

	desc.trunk.base_radius = {0.2f, 0.3f};
	desc.trunk.segment_length = {0.5f, 0.8f};
	desc.trunk.segment_count = [](float){
		return geom::interval{4, 7};
	};
	desc.trunk.lean_angle = {-15.f, 15.f};
	desc.trunk.rotation_angle = {-180.f, 180.f};
	desc.trunk.up_tendency = 0.f;
	desc.trunk.initial_branching_distance = {20.f, 30.f};
	desc.trunk.branching_distance = {0.1f, 0.15f};
	desc.trunk.branching_children_count = {1, 2};
	desc.trunk.branching_rotation = {60, 180.f};
	desc.trunk.branching_angle = {30.f, 75.f};
	desc.trunk.branching_size_multiplier = {1.f, 1.f};
	desc.trunk.splitting_probability = 0.1f;
	desc.trunk.splitting_angle = {15.f, 60.f};

	desc.branch.base_radius = {0.04f, 0.05f};
	desc.branch.segment_length = {0.1f, 0.15f};
	desc.branch.segment_count = [](float h){
		h = std::max(h / 4.f, 0.f);
//		h = std::min(h, 1.f);

//		h = 2.f * h - 1.f;

//		float t = std::sqrt(1.f - h * h);

		return geom::interval<int>{6, 10};
	};
	desc.branch.lean_angle = {-15.f, 15.f};
	desc.branch.rotation_angle = {-30.f, 30.f};
	desc.branch.up_tendency = 0.f;
	desc.branch.initial_branching_distance = {0.4f, 0.5f};
	desc.branch.branching_distance = {0.4f, 0.5f};
	desc.branch.branching_children_count = {1, 1};
	desc.branch.branching_rotation = {175.f, 185.f};
	desc.branch.branching_angle = {30.f, 60.f};
	desc.branch.branching_size_multiplier = {0.8f, 0.8f};
	desc.branch.splitting_probability = 0.1f;
	desc.branch.splitting_angle = {15.f, 60.f};

	return desc;
}

struct tree_description
{
	struct node
	{
		geom::point<float, 3> position;
		float radius;
	};

	struct branch
	{
		std::size_t parent;
		std::vector<node> nodes;
	};

	std::vector<branch> branches;
};

template <typename RNG>
tree_description generate(tree_species_description const & species, RNG && rng)
{
	// warm up!
	for (int i = 0; i < 16; ++i)
		rng();

	tree_description result;

	struct frame
	{
		geom::point<float, 3> pos;
		geom::vector<float, 3> z;
		geom::vector<float, 3> x;
	};

	auto generate_curve = util::recursive([&](auto & self, frame f, tree_species_description::curve_description const & curve_desc, tree_species_description::curve_description const & children_desc,
		int level, float min_radius, int max_segments, float h, float branching_distance, float branching_rotation) -> void
	{
		std::size_t id = result.branches.size();
		result.branches.emplace_back();

		float base_radius = random::uniform_distribution<float>{curve_desc.base_radius}(rng);
		base_radius = std::min(base_radius, min_radius);

		int segment_count = random::uniform_distribution<int>{curve_desc.segment_count(h)}(rng);
		segment_count = std::min(max_segments, segment_count);

		float expected_height = segment_count * curve_desc.segment_length.center();

		result.branches[id].parent = 0;
		result.branches[id].nodes.push_back({f.pos, base_radius});

		for (int i = 0; i < segment_count; ++i)
		{
			float t = (i + 1.f) / segment_count;

			float length = random::uniform_real_distribution<float>{curve_desc.segment_length}(rng);

			float lean = geom::rad(random::uniform_real_distribution<float>{curve_desc.lean_angle}(rng));

			f.z = geom::axis_rotation<float>{f.x, lean}(f.z);

			if (curve_desc.up_tendency > 0.f)
				f.z = geom::slerp(f.z, geom::vector{0.f, 0.f, 1.f}, t * curve_desc.up_tendency);
			else if (curve_desc.up_tendency < 0.f)
				f.z = geom::slerp(f.z, geom::vector{0.f, 0.f, -1.f}, - t * curve_desc.up_tendency);

			float rotation = geom::rad(random::uniform_real_distribution<float>{curve_desc.rotation_angle}(rng));

			f.x = geom::axis_rotation<float>{f.z, rotation}(f.x);

			auto new_pos = f.pos + f.z * length;

			float new_radius = (1.f - t) * base_radius;

			if (level < 3)
			{
				float available_branching_length = length;
				while (branching_distance < available_branching_length)
				{
					available_branching_length -= branching_distance;

					int children_count = random::uniform_distribution<int>{curve_desc.branching_children_count}(rng);

					for (int c = 0; c < children_count; ++c)
					{
						float angle = random::uniform_distribution<float>{curve_desc.branching_angle}(rng);

						float rotation = 0;
						if (children_count > 1)
							rotation = c * 2.f * geom::pi / children_count;

						float branch_t = 1.f - available_branching_length / length;

						geom::vector y{0.f, 0.f, 1.f};

						frame child_f;
						child_f.pos = geom::lerp(f.pos, new_pos, branch_t);
						child_f.z = geom::slerp(f.z, f.x, angle / 90.f);
						child_f.z = geom::axis_rotation<float>{f.z, branching_rotation + rotation}(child_f.z);
						child_f.x = geom::normalized(geom::cross(y, child_f.z));

						float multiplier = random::uniform_distribution<float>{curve_desc.branching_size_multiplier}(rng);

						float branching_distance = random::uniform_distribution<float>{children_desc.initial_branching_distance}(rng);

						self(child_f, children_desc, children_desc, level + 1,
							geom::lerp(result.branches[id].nodes.back().radius, new_radius, branch_t) * multiplier,
							(level == 0) ? 1024 : (segment_count - i) * multiplier,
							(level == 0) ? child_f.pos[2] / expected_height : h,
							branching_distance, 0.f);
					}

					branching_rotation += geom::rad(random::uniform_distribution<float>{curve_desc.branching_rotation}(rng));
					branching_distance = random::uniform_distribution<float>{curve_desc.branching_distance}(rng);
				}

				branching_distance -= available_branching_length;
			}

			f.pos = new_pos;

			result.branches[id].nodes.push_back({f.pos, new_radius});

			if (i + 1 < segment_count && random::uniform_distribution<float>{}(rng) < curve_desc.splitting_probability)
			{
				frame f1 = f;
				frame f2 = f;

				float a1 = geom::rad(random::uniform_distribution<float>{curve_desc.splitting_angle}(rng));
				float a2 = geom::rad(random::uniform_distribution<float>{curve_desc.splitting_angle}(rng));

				f1.z = geom::axis_rotation<float>{f.x, a1}(f1.z);
				f2.z = geom::axis_rotation<float>{f.x, -a2}(f2.z);

				float h = f.pos[2] / expected_height;

				self(f1, curve_desc, children_desc, level, new_radius, segment_count - i - 1, h, branching_distance, branching_rotation);
				self(f2, curve_desc, children_desc, level, new_radius, segment_count - i - 1, h, branching_distance, branching_rotation);

				break;
			}
		}
	});

	frame starting_frame;
	starting_frame.pos = {0.f, 0.f, 0.f};
	starting_frame.z = {0.f, 0.f, 1.f};
	starting_frame.x = {1.f, 0.f, 0.f};

	float initial_orientation = random::uniform_distribution<float>{0.f, 2.f * geom::pi}(rng);
	starting_frame.x = geom::axis_rotation<float>{starting_frame.z, initial_orientation}(starting_frame.x);

	generate_curve(starting_frame, species.trunk, species.branch, 0, std::numeric_limits<float>::infinity(), 1024, 0.f,
		random::uniform_distribution<float>{species.trunk.initial_branching_distance}(rng), 0.f);

	return result;
}

static char const vertex_source[] =
R"(#version 330

uniform mat4 u_transform;

layout (location = 0) in vec4 in_position;
layout (location = 1) in vec4 in_color;

out vec4 color;
out vec3 pos;

void main()
{
	gl_Position = u_transform * in_position;
	color = in_color;
	pos = in_position.xyz;
}
)";

static char const fragment_source[] =
R"(#version 330

in vec4 color;
in vec3 pos;

out vec4 out_color;

void main()
{
	vec3 n = normalize(cross(dFdx(pos), dFdy(pos)));

	vec3 light = normalize(vec3(1.0, 1.0, 1.0));

	float l = 0.5 + 0.5 * dot(n, light);

	out_color = vec4(l * color.rgb, 1.0);
}
)";

struct tree_app
	: app::app
{
	tree_species_description species = oak_species();

	tree_description tree_desc;

	gfx::mesh tree_mesh;

	gfx::program tree_program{vertex_source, fragment_source};

	geom::spherical_camera camera;

	std::uint64_t seed = 0;

	tree_app()
		: app("Tree")
	{
		tree_desc = generate(species, random::generator{seed, 0ull});
		update_mesh();

		setup_camera();
	}

	void update_mesh();
	void setup_camera();

	void on_mouse_move(int x, int y, int dx, int dy) override;

	void on_mouse_wheel(int delta) override;

	void on_resize(int width, int height) override;

	void on_key_down(SDL_Keycode key) override;

	void present() override;
};

struct vertex
{
	geom::point<float, 3> position;
	gfx::color_rgba color;
};

void tree_app::update_mesh()
{
	tree_mesh.setup<geom::point<float, 3>, gfx::normalized<gfx::color_rgba>>();

	std::vector<vertex> vertices;
	std::vector<std::uint32_t> indices;

	{
		int const basement_size = 5;
		float basement_offset = 0.05f;
		gfx::color_rgba basement_color { 127, 127, 127, 255 };

		for (int x = -basement_size; x < basement_size; ++x)
		{
			for (int y = -basement_size; y < basement_size; ++y)
			{
				auto base = vertices.size();

				vertices.push_back({{x + basement_offset, y + basement_offset, 0.f}, basement_color});
				vertices.push_back({{x + 1 - basement_offset, y + basement_offset, 0.f}, basement_color});
				vertices.push_back({{x + basement_offset, y + 1 - basement_offset, 0.f}, basement_color});
				vertices.push_back({{x + 1 - basement_offset, y + 1 - basement_offset, 0.f}, basement_color});

				indices.push_back(base + 0);
				indices.push_back(base + 1);
				indices.push_back(base + 2);
				indices.push_back(base + 2);
				indices.push_back(base + 1);
				indices.push_back(base + 3);
			}
		}
	}

	gfx::color_rgba trunk_color { 127, 63, 0, 255 };

	auto build_branch = [&](std::vector<tree_description::node> const & nodes)
	{
		int const N = 6;

		auto const base = vertices.size();

		auto z = geom::vector{0.f, 0.f, 1.f};
		auto x = geom::ort(z);

		for (std::size_t i = 0; i < nodes.size(); ++i)
		{
			if (i + 1 < nodes.size())
				z = geom::normalized(nodes[i + 1].position - nodes[i].position);
			else if (nodes.size() >= 2)
				z = geom::normalized(nodes[i].position - nodes[i - 1].position);

			auto y = geom::cross(z, x);

			geom::gram_schmidt(z, x, y);

			for (int k = 0; k < N; ++k)
			{
				float a = (k * geom::pi * 2.f) / N;

				vertices.push_back({nodes[i].position + (x * std::cos(a) + y * std::sin(a)) * nodes[i].radius, trunk_color});
			}
		}

		for (std::size_t i = 0; i + 1 < nodes.size(); ++i)
		{
			for (int k = 0; k < N; ++k)
			{
				int kk = (k + 1) % N;

				indices.push_back(base + i * N + k);
				indices.push_back(base + i * N + kk);
				indices.push_back(base + i * N + k + N);

				indices.push_back(base + i * N + k + N);
				indices.push_back(base + i * N + kk);
				indices.push_back(base + i * N + kk + N);
			}
		}
	};

	for (auto const & b : tree_desc.branches)
		build_branch(b.nodes);

	tree_mesh.load(vertices, indices, gl::TRIANGLES, gl::STATIC_DRAW);
}

void tree_app::setup_camera()
{
	camera.fov_y = geom::rad(45.f);
	camera.near_clip = 0.1f;
	camera.far_clip = 100.f;
	camera.target = {0.f, 0.f, 0.f};
	camera.elevation_angle = geom::rad(45.f);
	camera.azimuthal_angle = 0.f;
	camera.distance = 10.f;
}

void tree_app::on_mouse_move(int x, int y, int dx, int dy)
{
	app::on_mouse_move(x, y, dx, dy);

	if (is_middle_button_down())
	{
		camera.azimuthal_angle -= dx * 0.01f;
		camera.elevation_angle += dy * 0.01f;
	}

	if (is_right_button_down())
	{
		camera.target[2] += dy * 0.001f * camera.distance;
	}
}

void tree_app::on_mouse_wheel(int delta)
{
	camera.distance *= std::pow(0.8f, delta);
}

void tree_app::on_resize(int width, int height)
{
	app::on_resize(width, height);
	camera.set_fov(camera.fov_y, (1.f * width) / height);
}

void tree_app::on_key_down(SDL_Keycode key)
{
	if (key == SDLK_SPACE)
	{
		++seed;

		tree_desc = generate(species, random::generator{std::uint64_t{seed}, 0ull});
		update_mesh();
	}
}

void tree_app::present()
{
	gl::ClearColor(0.7f, 0.7f, 1.f, 1.f);
	gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

	gl::Enable(gl::DEPTH_TEST);
	gl::DepthFunc(gl::LEQUAL);

	gl::Enable(gl::CULL_FACE);

	tree_program.bind();
	tree_program["u_transform"] = camera.transform();
	tree_mesh.draw();
}

int main()
{
	return app::main<tree_app>();
}