#include <psemek/app/application_base.hpp>
#include <psemek/app/default_application_factory.hpp>
#include <psemek/math/box.hpp>
#include <psemek/math/mesh.hpp>
#include <psemek/math/intersection.hpp>
#include <psemek/math/camera.hpp>
#include <psemek/math/homogeneous.hpp>
#include <psemek/math/constants.hpp>
#include <psemek/math/math.hpp>
#include <psemek/math/gradient.hpp>
#include <psemek/gfx/gl.hpp>
#include <psemek/gfx/mesh.hpp>
#include <psemek/gfx/renderer/simple.hpp>
#include <psemek/gfx/program.hpp>
#include <psemek/gfx/texture.hpp>
#include <psemek/gfx/fullscreen.hpp>
#include <psemek/pcg/perlin.hpp>
#include <psemek/pcg/fractal.hpp>
#include <psemek/random/generator.hpp>
#include <psemek/random/device.hpp>
#include <psemek/random/uniform_sphere.hpp>
#include <psemek/util/range.hpp>
#include <psemek/util/clock.hpp>
#include <psemek/util/to_string.hpp>
#include <psemek/async/threadpool.hpp>
#include <psemek/log/log.hpp>

#include <iostream>

using namespace psemek;

template <typename Iterator>
auto barycenter(Iterator begin, Iterator end)
{
	auto p0 = *begin++;

	using vector_type = decltype(p0 - p0);

	auto v = vector_type::zero();

	std::size_t count = 1;

	for (; begin != end; ++begin, ++count)
	{
		v += (*begin - p0);
	}

	return p0 + v / static_cast<typename vector_type::scalar_type>(count);
}

template <typename Container>
auto barycenter(Container const & c)
{
	return barycenter(util::xbegin(c), util::xend(c));
}

std::vector<math::point<float, 3>> intersection(math::vector<float, 4> const & f, std::vector<math::point<float, 3>> const & vertices, std::vector<math::segment<std::uint32_t>> const & edges)
{
	std::vector<math::point<float, 3>> points;

	for (auto e : edges)
	{
		auto p0 = vertices[e[0]];
		auto p1 = vertices[e[1]];
		auto f0 = dot(f, math::homogeneous(p0));
		auto f1 = dot(f, math::homogeneous(p1));

		if ((f0 >= 0.f) ^ (f1 >= 0.f))
		{
			points.push_back(p0 - (p1 - p0) * f0 / (f1 - f0));
		}
	}

	math::vector<float, 3> const n { f[0], f[1], f[2] };

	if (!points.empty())
	{
		auto const & p0 = points[0];

		std::sort(points.begin() + 1, points.end(), [&](math::point<float, 3> const & p1, math::point<float, 3> const & p2){
			return math::dot(math::cross(p1 - p0, p2 - p0), n) > 0.f;
		});
	}

	return points;
}

template <typename Vertex, std::size_t N, typename Index = std::uint32_t>
struct mesh_builder
{
	std::vector<Vertex> vertices;
	std::vector<math::simplex<Index, N>> indices;

	void add(std::vector<Vertex> const & v, std::vector<math::simplex<Index, N>> const & i)
	{
		Index const base = static_cast<Index>(vertices.size());
		std::size_t begin = indices.size();

		vertices.insert(vertices.end(), v.begin(), v.end());
		indices.insert(indices.end(), i.begin(), i.end());

		for (std::size_t i = begin; i < indices.size(); ++i)
		{
			for (auto & idx : indices[i].points)
				idx += base;
		}
	}
};

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

uniform mat4 u_transform;

layout (location = 0) in vec4 in_position;
layout (location = 1) in vec3 in_texcoord;

out vec3 texcoord;

void main()
{
	gl_Position = u_transform * in_position;
	texcoord = in_texcoord;
}
)";

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

uniform sampler3D u_texture;
uniform sampler3D u_shadow;
uniform float u_step;
uniform float u_max_density;
uniform float u_min_harmonic;
uniform float u_max_harmonic;
uniform vec3 u_light;

in vec3 texcoord;
out vec4 out_color;

void main()
{
	float pi = 3.1415926535;

	float a = texture(u_texture, texcoord).r * u_step * u_max_density;
	vec4 sv = texture(u_shadow, texcoord);

	sv = vec4(1.0, 1.0, 1.0, 1.0) * u_min_harmonic + sv * (u_max_harmonic - u_min_harmonic);

	float g = sqrt(3.0 / 4.0 / pi);

	float s = sv.w / sqrt(2.0 * pi) + dot(sv.xyz, u_light) * g;

	float l = 1.0 - s;
	l = max(0.0, l);

	float ambient = 0.0;

	vec3 color = vec3(1.0, 1.0, 1.0) * (ambient + (1.0 - ambient) * l);
	out_color = vec4(color, a);
}
)";

struct cloud_app
	: app::application_base
{
	math::spherical_camera camera;

	float step;

	float max_density = 2.f;
	math::interval<float> harmonic_range;

	util::clock<std::chrono::duration<float>> clock;

	math::box<float, 3> bbox;
	std::vector<math::point<float, 3>> bbox_vertices;
	std::vector<math::segment<std::uint32_t>> bbox_edges;

	gfx::mesh bbox_mesh;

	gfx::mesh slice_mesh;

	gfx::mesh light_mesh;
	std::vector<math::vector<float, 3>> dirs;

	gfx::simple_renderer renderer;

	gfx::program cloud_program{cloud_vertex_source, cloud_fragment_source};

	gfx::texture_3d cloud_texture;
	gfx::texture_3d shadow_texture;

	cloud_app(std::size_t size)
	{
		math::vector<int, 3> cloud_size{2 * size, size, size};
		bbox = {{{-2.f, 2.f}, {-1.f, 1.f}, {-1.f, 1.f}}};

		step = bbox[0].length() / cloud_size[0];

		camera.fov_y = math::rad(45.f);
		camera.near_clip = 0.01f;
		camera.far_clip = 100.f;

		camera.target = bbox.center();
		camera.azimuthal_angle = 0.f;
		camera.elevation_angle = 0.f;
		camera.distance = 10.f;

		bbox_vertices = math::vertices(bbox);
		bbox_edges = math::edges(bbox);

		bbox_mesh.setup<math::point<float, 3>>();
		bbox_mesh.load(bbox_vertices, bbox_edges);

		slice_mesh.setup<math::point<float, 3>, math::vector<float, 3>>();

		light_mesh.setup<math::point<float, 3>>();

		{
			async::threadpool bg("bg");

			random::generator rng(random::device{});
			random::uniform_sphere_vector_distribution<float, 3> d;

			std::vector<util::ndarray<math::vector<float, 3>, 3>> grad(4);
			std::vector<float> weights(grad.size());
			float weight_sum = 0.f;

			for (std::size_t o = 0; o < grad.size(); ++o)
			{
				grad[o].resize({(16 << o), (8 << o), (8 << o)});
				for (auto & v : grad[o])
				{
					v = d(rng);
				}
				weights[o] = std::pow(1.f / (1 << o), 1.f);
				weight_sum += weights[o];
			}

			for (float & w : weights)
				w /= weight_sum;

			pcg::fractal<pcg::perlin<float, 3>> perlin(std::move(grad), std::move(weights));

			// can't use template argument deduction for first argument due to gcc bug
			// https://gcc.gnu.org/bugzilla/show_bug.cgi?id=89062
			math::gradient<float> g(std::make_pair(0.2f, 0.f), math::easing_type::quadratic_out, std::pair{0.3f, max_density});

			util::ndarray<std::uint8_t, 3> cloud_data({cloud_size[0], cloud_size[1], cloud_size[2]});

			for (std::size_t z = 0; z < cloud_data.depth(); ++z)
			{
				for (std::size_t y = 0; y < cloud_data.height(); ++y)
				{
					bg.post([&, z, y]{
						for (std::size_t x = 0; x < cloud_data.width(); ++x)
						{
							math::vector<float, 3> p;
							p[0] = (x * 1.f) / cloud_data.width();
							p[1] = (y * 1.f) / cloud_data.height();
							p[2] = (z * 1.f) / cloud_data.depth();

							float v = perlin(p);

							auto d = p - math::vector{0.5f, 0.5f, 0.5f};
							v *= std::max(0.f, 1.f - math::length(d) / 0.5f);

							cloud_data(x, y, z) = math::clamp(g(v) / max_density * 255.f, {0.f, 255.f});
						}
					});
				}
			}

			bg.wait();
			log::info() << "Generated a cloud";

			cloud_texture.load(cloud_data);
			cloud_texture.repeat();
			cloud_texture.linear_filter();
			cloud_texture.generate_mipmap();

			auto value_at = [this, &cloud_data](math::point<float, 3> const & p)
			{
				math::vector<float, 3> d;
				d[0] = (p[0] - bbox[0].min) / bbox[0].length() * cloud_data.width();
				d[1] = (p[1] - bbox[1].min) / bbox[1].length() * cloud_data.height();
				d[2] = (p[2] - bbox[2].min) / bbox[2].length() * cloud_data.depth();

				d[0] -= 0.5f;
				d[1] -= 0.5f;
				d[2] -= 0.5f;

				if (d[0] < 0.0f) return 0.f;
				if (d[0] >= cloud_data.width() - 1) return 0.f;
				if (d[1] < 0.0f) return 0.f;
				if (d[1] >= cloud_data.height() - 1) return 0.f;
				if (d[2] < 0.0f) return 0.f;
				if (d[2] >= cloud_data.depth() - 1) return 0.f;

				math::vector<int, 3> i;
				i[0] = std::floor(d[0]);
				i[1] = std::floor(d[1]);
				i[2] = std::floor(d[2]);

				math::vector<float, 3> t;
				t[0] = d[0] - i[0];
				t[1] = d[1] - i[1];
				t[2] = d[2] - i[2];

				float v000 = cloud_data(i[0], i[1], i[2]);
				float v100 = cloud_data(i[0] + 1, i[1], i[2]);
				float v010 = cloud_data(i[0], i[1] + 1, i[2]);
				float v110 = cloud_data(i[0] + 1, i[1] + 1, i[2]);
				float v001 = cloud_data(i[0], i[1], i[2] + 1);
				float v101 = cloud_data(i[0] + 1, i[1], i[2] + 1);
				float v011 = cloud_data(i[0], i[1] + 1, i[2] + 1);
				float v111 = cloud_data(i[0] + 1, i[1] + 1, i[2] + 1);

				float v00 = math::lerp(v000, v100, t[0]);
				float v10 = math::lerp(v010, v110, t[0]);
				float v01 = math::lerp(v001, v101, t[0]);
				float v11 = math::lerp(v011, v111, t[0]);

				float v0 = math::lerp(v00, v10, t[1]);
				float v1 = math::lerp(v01, v11, t[1]);

				return math::lerp(v0, v1, t[2]) / 255.f * max_density;
			};

			util::ndarray<math::vector<float, 4>, 3> cloud_shadow_f(cloud_data.dims());

			dirs.resize(32);

			// Fibonacci spiral:
			float const phi = (1.f + std::sqrt(5.f)) / 2.f;
			for (int i = 0; i < dirs.size(); ++i)
			{
				float a = 2.f * math::pi * (i / phi);
				float b = std::acos(1.f - (2.f * i) / dirs.size());

				dirs[i][0] = std::cos(a) * std::sin(b);
				dirs[i][1] = std::sin(a) * std::sin(b);
				dirs[i][2] = std::cos(b);
			}

			float diag = std::sqrt(math::sqr(bbox[0].length()) + math::sqr(bbox[1].length()) + math::sqr(bbox[2].length()));
			int steps = diag / step;

			auto process = [&, this](auto const & idx)
			{
				float const step = bbox[0].length() / cloud_size[0];

				math::point<float, 3> o;
				for (std::size_t i = 0; i < 3; ++i)
					o[i] = bbox[i].min + bbox[i].length() * (idx[i] + 0.5f) / cloud_shadow_f.dim(i);

				float sum_0 = 0.f;
				float sum_x = 0.f;
				float sum_y = 0.f;
				float sum_z = 0.f;
				int count = 0;
				for (auto const & dir : dirs)
				{
					float density = 0.f;

					math::point<float, 3> p = o + dir * (steps * step);

					for (int s = steps; s > 0; --s)
					{
						float a = value_at(p) * step;
						density = a + density * (1.f - a);

						p -= dir * step;
					}

					sum_0 += density * 2.f * std::sqrt(math::pi);
					sum_x += density * dir[0] * std::sqrt(12.f * math::pi);
					sum_y += density * dir[1] * std::sqrt(12.f * math::pi);
					sum_z += density * dir[2] * std::sqrt(12.f * math::pi);
					++count;
				}

				float value_0 = sum_0 / count;
				float value_x = sum_x / count;
				float value_y = sum_y / count;
				float value_z = sum_z / count;

				cloud_shadow_f(idx)[0] = value_x;
				cloud_shadow_f(idx)[1] = value_y;
				cloud_shadow_f(idx)[2] = value_z;
				cloud_shadow_f(idx)[3] = value_0;
			};

			std::size_t total = cloud_data.size();
			std::atomic<std::size_t> done = 0;

			for (std::size_t z = 0; z < cloud_data.depth(); ++z)
			{
				for (std::size_t y = 0; y < cloud_data.height(); ++y)
				{
					bg.dispatch([&, z, y]{
						for (std::size_t x = 0; x < cloud_data.width(); ++x)
						{
							process(std::array<std::size_t, 3>{{x, y, z}});
						}
						done += cloud_data.width();
						log::info() << (done.load() * 100.f / total) << " %";
					});
				}
			}

			bg.wait();

			log::info() << "Finished!";

			math::interval<float> range_0;
			math::interval<float> range_d;

			for (auto const & v : cloud_shadow_f)
			{
				for (std::size_t i = 0; i < 4; ++i)
				{
					harmonic_range |= v[i];
				}

				range_0 |= v[3];
				range_d |= v[0];
				range_d |= v[1];
				range_d |= v[2];
			}

			log::info() << "Full: " << harmonic_range;

			log::info() << "0: " << range_0;
			log::info() << "XYZ: " << range_d;

			log::info() << "V0 in center: " << cloud_shadow_f(cloud_data.width() / 2, cloud_data.height() / 2, cloud_data.depth() / 2)[3];

			for(std::size_t z = 0; z < 2; ++z)
			{
				for(std::size_t y = 0; y < 2; ++y)
				{
					for(std::size_t x = 0; x < 2; ++x)
					{
						log::info() << "V0 in corner " << x << y << z << " : " << cloud_shadow_f(x*(cloud_data.width() - 1), y*(cloud_data.height()-1), z*(cloud_data.depth()-1))[3];
					}
				}
			}

			util::ndarray<math::vector<std::uint8_t, 4>, 3> cloud_shadow(cloud_data.dims());

			for (auto const & idx : cloud_shadow.indices())
			{
				for (std::size_t i = 0; i < 4; ++i)
				{
					float v = (cloud_shadow_f(idx)[i] - harmonic_range.min) / harmonic_range.length();
					cloud_shadow(idx)[i] = math::clamp(255.f * v, {0.f, 255.f});
				}
			}

			shadow_texture.load(cloud_shadow);
			shadow_texture.clamp();
			shadow_texture.linear_filter();
			shadow_texture.generate_mipmap();
		}
	}

	void on_event(app::resize_event const & event) override
	{
		app::application_base::on_event(event);

		gl::Viewport(0, 0, event.size[0], event.size[1]);

		float aspect_ratio = (1.f * event.size[0]) / event.size[1];
		camera.set_fov(camera.fov_y, aspect_ratio);
	}

	void on_event(app::mouse_move_event const & event) override
	{
		auto const old_mouse = state().mouse;

		app::application_base::on_event(event);

		if (state().mouse_button_down.contains(app::mouse_button::middle))
		{
			auto const delta = event.position - old_mouse;
			camera.azimuthal_angle -= delta[0] * 0.01f;
			camera.elevation_angle += delta[1] * 0.01f;
		}
	}

	void on_event(app::mouse_wheel_event const & event) override
	{
		app::application_base::on_event(event);

		camera.distance *= std::pow(0.8f, event.delta);
	}

	void update_slice_mesh()
	{
		auto n = -camera.direction();
		math::vector<float, 4> f {n[0], n[1], n[2], 0.f};

		math::interval<float> range;

		for (auto p : bbox_vertices)
		{
			range |= math::dot(f, math::homogeneous(p));
		}

		int count = std::ceil(range.length() / step);

		struct vertex
		{
			math::point<float, 3> pos;
			math::vector<float, 3> tc;
		};

		mesh_builder<vertex, 2> builder;

		for (int s = count - 1; s > 0; --s)
		{
			f[3] = - (range.max - (range.length() * s) / count);

			auto slice_vertices = intersection(f, bbox_vertices, bbox_edges);
			auto slice_indices = math::triangulate_convex(slice_vertices);

			std::vector<vertex> vertices;
			for (auto p : slice_vertices)
			{
				vertex & v = vertices.emplace_back();
				v.pos = p;
				for (std::size_t i = 0; i < 3; ++i)
					v.tc[i] = (p[i] - bbox[i].min) / bbox[i].length();
			}

			builder.add(vertices, slice_indices);
		}

		slice_mesh.load(builder.vertices, builder.indices);
	}

	void update() override
	{}

	void present() override
	{
		update_slice_mesh();

		float t = clock.count();

		math::vector<float, 3> light_dir = math::normalized(math::vector{std::cos(t), std::sin(t), 0.5f});

		{
			std::vector<math::point<float, 3>> light_vertices;
			auto o = math::point<float, 3>::zero() + light_dir * 4.f;
			float s = 0.2f;
			for (auto d : dirs)
			{
				light_vertices.push_back(o);
				light_vertices.push_back(o + d * s);
			}
			light_mesh.load(light_vertices, gl::LINES);
		}

		gl::Enable(gl::CULL_FACE);

		gl::ClearColor(0.5f, 0.5f, 1.f, 1.f);
		gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

		gl::Disable(gl::BLEND);

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

		renderer.push({&light_mesh, gfx::yellow.as_color_rgba()});
		renderer.push({&bbox_mesh, gfx::red.as_color_rgba()});
		renderer.render({camera.transform()});

		gl::Enable(gl::BLEND);
		gl::BlendFunc(gl::SRC_ALPHA, gl::ONE_MINUS_SRC_ALPHA);

		gl::DepthMask(gl::FALSE);

		cloud_program.bind();
		cloud_program["u_transform"] = camera.transform();
		cloud_program["u_texture"] = 0;
		cloud_program["u_shadow"] = 1;
		cloud_program["u_step"] = step;
		cloud_program["u_max_density"] = max_density;
		cloud_program["u_min_harmonic"] = harmonic_range.min;
		cloud_program["u_max_harmonic"] = harmonic_range.max;
		cloud_program["u_light"] = math::normalized(light_dir);
		gl::ActiveTexture(gl::TEXTURE0);
		cloud_texture.bind();
		gl::ActiveTexture(gl::TEXTURE1);
		shadow_texture.bind();
		slice_mesh.draw();

		gl::DepthMask(gl::TRUE);
	}
};

namespace psemek::app
{

	std::unique_ptr<application::factory> make_application_factory()
	{
		return default_application_factory({.name = "Cloud example"}, [](application::options const &, application::context const & context)
		-> std::unique_ptr<application> {
			if (context.args.size() != 1 && context.args.size() != 2)
			{
				std::cout << "Usage: " << context.args[0] << " [ size ]\n";
				return nullptr;
			}

			std::size_t size = 32;
			if (context.args.size() == 2)
				size = util::from_string<std::size_t>(context.args[1]);
			return std::make_unique<cloud_app>(size);
		});
	}

}
