#pragma once

#include "utilities.hpp"

#include <cassert>
#include <utility>
#include <array>
#include <vector>
#include <algorithm>
#include <iostream>
#include <functional>

// Position {0,0} is bottom left.
// Position {W-1, H-1} is top right.
// template <typename T>
struct DynamicGrid{
	using Position = std::pair<int, int>;

	template <typename Input>
	DynamicGrid(int W, int H, Input it)
	: W(W)
	, H(H)
	, grid(W*H, 0)
	{
		std::copy_n(it, W*H, std::begin(grid));
	}

	//! \brief Returns true if p represents a valid position
	auto valid(Position const & p) const {
		return p.first >= 0 && p.second >= 0
		    && p.first <  W && p.second < H;
	}

	//! \brief Returns true if p represents an empty cell
	auto empty(Position const & p) const {
		assert(valid(p));
		return get(p) == 0;
	}

	//! \brief Returns all neighbors of p
	//! Only returns valid neighbors (and p does not need to be valid)
	auto neighbors(Position const & p) const {
		small_vector<Position, 4> ret;

		static Position nbs[] = {
			{-1, 0}, {1, 0}, {0, 1}, {0, -1}
		};

		for(auto&& n : nbs){
			auto q = Position{p.first + n.first, p.second + n.second};
			if(valid(q)){
				ret.push_back(q);
			}
		}

		return ret;
	}

	//! \brief Get (mutable) value at p
	int & get(Position const & p){
		assert(valid(p));
		return grid[static_cast<size_t>(p.first + p.second*W)];
	}

	//! \brief Get value at p
	int const& get(Position const & p) const {
		assert(valid(p));
		return grid[static_cast<size_t>(p.first + p.second*W)];
	}

	//! \brief Pretty prints grid to out
	void print(std::ostream& out) const {
		for(auto y = H; y; --y){
			for(auto x = 0; x < W; ++x){
				out << get({x, y-1}) << (x == W-1 ? '\n' : ' ');
			}
		}
	}

private:
	int W;
	int H;
	std::vector<int> grid;

	//! \brief Returns true if the whole column at x is empty
	auto empty_column(int x){
		for(auto y = 0; y < H; ++y){
			if(!empty({x, y})){
				return false;
			}
		}
		return true;
	}

	//! \brief Helper function for all_positions()
	static auto all_positions_impl(int W, int H) {
		std::vector<Position> r;
		for(int y = 0; y < H; ++y)
			for(int x = 0; x < W; ++x)
				r.emplace_back(x, y);
		return r;
	}

public:
	//! \brief Returns an array of all valid positions
	auto const & all_positions() const {
		static auto v = all_positions_impl(W, H);
		return v;
	}

private:
	//! \brief Single vertical collapsing step
	auto vcollapse(){
		using namespace std;
		bool some_change = false;
		for(auto&& p : all_positions()){
			if(empty(p)) continue;

			auto q = p;
			q.second--;
			if(!valid(q)) continue;
			if(!empty(q)) continue;

			swap(get(p), get(q));
			some_change = true;
		}

		return some_change;
	}

	//! \brief Single horizontal collapsing step
	auto hcollapse(){
		using namespace std;
		bool some_change = false;
		for(auto x = 0; x < W-1; ++x){
			if(!empty_column(x)) continue;

			auto x2 = x+1;
			for(; x2 < W; ++x2){
				if(!empty_column(x2)) break;
			}
			if(x2 == W) return some_change;

			for(auto y = 0; y < H; ++y){
				swap(get({x, y}), get({x2, y}));
			}
			some_change = true;
		}
		return some_change;
	}

public:
	//! \brief Let the block fall (all the way)
	auto collapse(){
		while(vcollapse());
		while(hcollapse());
	}
};

template <typename URNG>
auto random_dynamic_grid(int W, int H, URNG&& r){
	std::uniform_int_distribution<int> dis(1, 2);
	std::vector<int> v(W*H);
	std::generate_n(std::begin(v), W*H, [&]{ return dis(r); });
	return DynamicGrid(W, H, std::begin(v));
}