#include <psemek/app/app.hpp>
#include <psemek/app/main.hpp>
#include <psemek/gfx/painter.hpp>
#include <psemek/gfx/gl.hpp>
#include <psemek/geom/scale.hpp>
#include <psemek/geom/camera.hpp>
#include <psemek/geom/constants.hpp>
#include <psemek/util/clock.hpp>

#include <random>

/*
No optimizations, 125:
[2020 Aug 29  09:54:27.348][ main][   info] Avg forces time:    0.000641577
[2020 Aug 29  09:54:27.348][ main][   info] Avg collision time: 0.000789326

No optimizations, 250:
[2020 Aug 29  09:49:48.476][ main][   info] Avg forces time:    0.00257929
[2020 Aug 29  09:49:48.476][ main][   info] Avg collision time: 0.00314664

No optimizations, 500:
[2020 Aug 29  09:44:02.200][ main][   info] Avg forces time:    0.00299429
[2020 Aug 29  09:44:02.200][ main][   info] Avg collision time: 0.00362124

No optimizations, 1000:
[2020 Aug 29  09:42:34.143][ main][   info] Avg forces time:    0.0072538
[2020 Aug 29  09:42:34.143][ main][   info] Avg collision time: 0.00873274

No optimizations, 2000:
[2020 Aug 29  09:44:35.334][ main][   info] Avg forces time:    0.029108
[2020 Aug 29  09:44:35.334][ main][   info] Avg collision time: 0.0345713

No optimizations, 4000:
[2020 Aug 29  09:50:15.647][ main][   info] Avg forces time:    0.118145
[2020 Aug 29  09:50:15.647][ main][   info] Avg collision time: 0.1402





*/

using namespace psemek;

struct particle
{
	geom::point<float, 2> pos;
	geom::vector<float, 2> vel;

	float radius;
	float mass;

	float T = 0.f;
};

struct myapp : app::app
{
	myapp()
		: app("Test app")
		, rng_{std::random_device{}()}
	{
		gl::ClearColor(1.f, 1.f, 1.f, 1.f);


		std::uniform_real_distribution<float> d{-200.f, 200.f};
		std::uniform_real_distribution<float> rr{0.5f, 2.f};

		for (int i = 0; i < 1000; ++i)
		{
			float r = rr(rng_);
			float m = r * r * r;

			float v = 0.1f;

			geom::point<float, 2> p;
			while (true)
			{
				p = {((i % 2) ? -300.f : 300.f) + d(rng_), d(rng_)};

				if (std::all_of(particles_.begin(), particles_.end(), [&](particle const & q){ return geom::distance(q.pos, p) > q.radius + r; }))
					break;
			}
			particles_.push_back({p, {v * d(rng_), v * d(rng_)}, r, m});
//			particles_.push_back({{ 400.f + d(rng), d(rng)}, {0.f, 0.f}, 1.f, 1.f});

//			particles_.push_back({{d(rng),  d(rng)}, {0.f, 0.f}, 1.f, 1.f});
//			particles_.push_back({{d(rng),  d(rng)}, {0.f, 0.f}, 1.f, 1.f});
		}
	}

	void on_resize(int width, int height) override
	{
		gl::Viewport(0, 0, width, height);

		window_size_ = {width, height};

		camera_ratio_ = static_cast<float>(width) / height;
	}

	void on_mouse_wheel(int delta) override
	{
		camera_size_ *= std::pow(0.8f, delta);
	}

	void on_left_button_down() override
	{
		if (mouse_)
		{
			geom::scale<float, 2> const flip_y({1.f, -1.f});

			float const scale = camera_size_ / window_size_[1];

			geom::point<int, 2> const screen_center { window_size_[0] / 2, window_size_[1] / 2 };
			auto target = camera_center_ + flip_y(geom::cast<float>(*mouse_ - screen_center) * scale);

			geom::point<float, 2> pos{100.f, 0.f};

			geom::vector<float, 2> vel = geom::normalized(target - pos) * 500.f;

			particles_.push_back({pos, vel, 1.f, 1.f});
		}
	}

	void on_right_button_down() override
	{
		camera_drag_ = mouse_;
	}

	void on_right_button_up() override
	{
		camera_drag_ = std::nullopt;
	}

	void on_mouse_move(int x, int y, int, int) override
	{
		mouse_ = {x, y};

		if (camera_drag_)
		{
			geom::scale<float, 2> const flip_y({1.f, -1.f});

			float const scale = camera_size_ / window_size_[1];

			camera_center_ += flip_y(geom::cast<float>(*camera_drag_ - *mouse_) * scale);
			camera_drag_ = mouse_;
		}
	}

	void update() override
	{
//		std::shuffle(particles_.begin(), particles_.end(), rng_);

		float const dt = 0.01f;
		float const G = 100.f;

		for (std::size_t step = 0; step < 1; ++step)
		{
			{
				util::clock<> clock;

				for (std::size_t i = 0; i < particles_.size(); ++i)
				{
					for (std::size_t j = i + 1; j < particles_.size(); ++j)
					{
						auto const r = particles_[i].pos - particles_[j].pos;

						float const l = length(r);

						auto const f = G * particles_[i].mass * particles_[j].mass * r / std::pow(l, 3.f);

						particles_[i].vel -= f * dt / particles_[i].mass;
						particles_[j].vel += f * dt / particles_[j].mass;
					}
				}

				total_forces_ += clock.count();
			}

			for (auto & p : particles_)
			{
				p.pos += p.vel * dt;
			}

			{
				util::clock<> clock;

				for (std::size_t iteration = 0; iteration < 10; ++iteration)
				{
					for (std::size_t i = 0; i < particles_.size(); ++i)
					{
						for (std::size_t j = i + 1; j < particles_.size(); ++j)
						{
							auto const r = particles_[i].pos - particles_[j].pos;

							float const l = length(r);

							float const R = particles_[i].radius + particles_[j].radius;

							auto const n = r / l;
/*
							if (l < R)
							{
								auto f = E * n * std::pow(l / R, -2.f);
								particles_[i].vel += f * dt / particles_[i].mass;
								particles_[j].vel -= f * dt / particles_[j].mass;

								auto const vij = particles_[i].vel - particles_[j].vel;
								auto dp = - K * dt * n * dot(n, vij) * 2.f / (1.f / particles_[i].mass + 1.f / particles_[j].mass);

								float Ei0 = geom::length_sqr(particles_[i].vel) * particles_[i].mass * 0.5f;
								particles_[i].vel += dp / particles_[i].mass;
								float Ei1 = geom::length_sqr(particles_[i].vel) * particles_[i].mass * 0.5f;
								particles_[i].T += Ei0 - Ei1;

								float Ej0 = geom::length_sqr(particles_[j].vel) * particles_[j].mass * 0.5f;
								particles_[j].vel -= dp / particles_[j].mass;
								float Ej1 = geom::length_sqr(particles_[j].vel) * particles_[j].mass * 0.5f;
								particles_[j].T += Ej0 - Ej1;

								float dT = particles_[i].T - particles_[j].T;
								particles_[i].T -= C * dT * dt;
								particles_[j].T += C * dT * dt;
							}
*/

							auto const vij = particles_[i].vel - particles_[j].vel;

							if (l < R && dot(vij, n) < 0.f)
							{
								auto const d = n * (R - l) / 2.f * 0.5f;// * (particles_[i].mass + particles_[j].mass);
								particles_[i].pos += d;
								particles_[j].pos -= d;

								auto dp = - n * dot(n, vij) * 2.f / (1.f / particles_[i].mass + 1.f / particles_[j].mass);

								particles_[i].vel += dp / particles_[i].mass;
								particles_[j].vel -= dp / particles_[j].mass;
							}
						}
					}
				}

				total_collisions_ += clock.count();
			}

			float Ep = 0.f;
			float Ek = 0.f;

			for (std::size_t i = 0; i < particles_.size(); ++i)
			{
				Ek += geom::length_sqr(particles_[i].vel) * particles_[i].mass / 2.f;

				for (std::size_t j = i + 1; j < particles_.size(); ++j)
				{
					Ep -= G * particles_[i].mass * particles_[j].mass / distance(particles_[i].pos, particles_[j].pos);
				}
			}

			log::info() << "Energy: " << Ek << " - " << (-Ep) << " = " << (Ek + Ep);

			++frame_count_;
		}
	}

	void present() override
	{
		gl::Clear(gl::COLOR_BUFFER_BIT);

		for (auto & p : particles_)
		{
			float c = 2.f / (std::exp(-p.T / 100000.f) + 1.f) - 1.f;
			painter_.circle(p.pos, p.radius, {static_cast<std::uint8_t>(c * 255.f), 0, 0, 255});
		}

		geom::orthographic_camera camera;
		camera.box[0].min = camera_center_[0] - camera_size_ * camera_ratio_ / 2.f;
		camera.box[0].max = camera_center_[0] + camera_size_ * camera_ratio_ / 2.f;
		camera.box[1].min = camera_center_[1] - camera_size_ / 2.f;
		camera.box[1].max = camera_center_[1] + camera_size_ / 2.f;
		camera.box[2].min = -1.f;
		camera.box[2].max =  1.f;

		painter_.render(camera.transform());
	}

	~myapp()
	{
		log::info() << "Avg forces time:    " << (total_forces_ / frame_count_);
		log::info() << "Avg collision time: " << (total_collisions_ / frame_count_);
	}

private:
	std::default_random_engine rng_;

	geom::vector<int, 2> window_size_;

	geom::point<float, 2> camera_center_ { 0.f, 0.f };
	float camera_size_ = 500.f;
	float camera_ratio_ = 1.f;

	std::optional<geom::point<int, 2>> mouse_;
	std::optional<geom::point<int, 2>> camera_drag_;

	gfx::painter painter_;

	std::vector<particle> particles_;

	int frame_count_ = 0;
	float total_forces_ = 0.f;
	float total_collisions_ = 0.f;
};

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