#include "lattice.h"

#include "LLBM/collide.h"
#include "LLBM/initialize.h"
#include "LLBM/bounce_back.h"
#include "LLBM/bounce_back_moving_wall.h"

#include <chrono>
#include <iostream>

#include "pattern/all.h"

using T = SWEEPLB_PRECISION;
using PATTERN = pattern::SWEEPLB_PATTERN<T>;

void simulate(Cuboid cuboid, std::size_t nStep) {
	const int nThread = omp_get_max_threads();

	Lattice<PATTERN> lattice(cuboid);

	LatticeMask<T> bulk_mask(cuboid.volume());
	LatticeMask<T> box_mask(cuboid.volume());
	LatticeMask<T> lid_mask(cuboid.volume());

	cuboid.traverse([&](int iX, int iY, int iZ, std::size_t iCell) {
		if (iZ == cuboid[2]-1) {
			lid_mask.set(iCell, true);
		} else if (iX == 0 || iX == cuboid[0]-1
			      || iY == 0 || iY == cuboid[1]-1
			      || iZ == 0) {
			box_mask.set(iCell, true);
		} else {
			bulk_mask.set(iCell, true);
		}
	});

	bulk_mask.serialize();
	box_mask.serialize();
	lid_mask.serialize();

	T tau = 0.51;
	T u_lid[] { 0.05, 0., 0. };

	for (std::size_t iStep = 0; iStep < 100; ++iStep) {
		lattice.apply(Operator(BgkCollideO(), bulk_mask, tau),
		              Operator(BounceBackO(), box_mask),
		              Operator(BounceBackMovingWallO(), lid_mask, u_lid));
		lattice.stream();
	}

	auto start = std::chrono::steady_clock::now();

	for (std::size_t iStep = 0; iStep < nStep; ++iStep) {
		lattice.apply(Operator(BgkCollideO(), bulk_mask, tau),
		              Operator(BounceBackO(), box_mask),
		              Operator(BounceBackMovingWallO(), lid_mask, u_lid));
		lattice.stream();
	}

	auto duration = std::chrono::duration_cast<std::chrono::duration<double>>(
		std::chrono::steady_clock::now() - start);

	std::cout << cuboid[0]
	          << ", " << nStep
	          << ", " << nThread
	          << ", " << (nStep * lattice.volume()) / (1e6 * duration.count())
	          << std::endl;

	lattice.write_momenta(bulk_mask, "result.vtk");
}

int main(int argc, char* argv[]) {
	const std::size_t n     = atoi(argv[1]);
	const std::size_t steps = atoi(argv[2]);
	simulate({ n, n, n}, steps);
}