#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
#include <tuple>
#include <cstdint>
#include <cstring>
#include <cassert>
#include "../lodepng.h"
#include "compute_host.h"
#include "defs.h"
#include "kernel.h"

using namespace std;


static void writeCounts(int W, int H, const vector<int> &counts, const char *fname) {
	ofstream f(fname);
	f << W << ' ' << H << '\n';
	for (int y = 0; y < H; y++) {
		for (int x = 0; x < W; x++) {
			if (x != 0) f << ' ';
			f << counts[W * y + x];
		}
		f << '\n';
	}
}

static tuple<int, int, vector<int>> readCounts(const char *fname) {
	ifstream f(fname);
	int W, H;
	f >> W >> H;
	vector<int> counts(W * H);
	for (int &v : counts) f >> v;
	return make_tuple(W, H, counts);
}

static int rankCounts(vector<int> &counts) {
	int maxcount = 0;
	for (int i = 0; i < (int)counts.size(); i++) {
		maxcount = max(maxcount, counts[i]);
	}

	vector<int> cumul(maxcount + 1, 0);
	for (int v : counts) cumul[v]++;
	cumul[0] = 0;
	for (int i = 1; i < (int)cumul.size(); i++) cumul[i] += cumul[i-1];
	// assert(cumul[maxcount + 1] == (int)counts.size());

	for (int &v : counts) v = cumul[v];

	return cumul[maxcount];
}

static vector<uint8_t> drawImage(int W, int H, const vector<int> &counts, int maxcount) {
	vector<uint8_t> image(3 * W * H);

	for (int y = 0; y < H; y++) {
		for (int x = 0; x < W; x++) {
			double value = (double)counts[W * y + x] / maxcount * 255;
			image[3 * (W * y + x) + 0] = value;
			image[3 * (W * y + x) + 1] = value;
			image[3 * (W * y + x) + 2] = value;
		}
	}

	return image;
}

int main(int argc, char **argv) {
	int W, H;
	vector<int> counts;

	if (argc <= 1) {
		W = H = 900;
		// Use 2.5 everywhere for Christensen
		const Com bottomLeft = Com(-1.5, -1.5);
		const Com topRight = Com(1.5, 1.5);

		// counts = computeHost(W, H, bottomLeft, topRight);
		Kernel().run_chunked(counts, W, H, bottomLeft, topRight, 42, 1 << 14);
		// Kernel().run_all(counts, W, H, bottomLeft, topRight, 42);

		writeCounts(W, H, counts, "out.txt");
	} else if (argc == 2 && strcmp(argv[1], "-h") != 0) {
		tie(W, H, counts) = readCounts(argv[1]);
	} else {
		cerr << "Usage: " << argv[0] << "            -- compute and draw" << endl;
		cerr << "Usage: " << argv[0] << " <out.txt>  -- draw already-computed data" << endl;
		return 1;
	}

	int maxcount = rankCounts(counts);

	vector<uint8_t> image = drawImage(W, H, counts, maxcount);

	assert(lodepng_encode24_file("out.png", image.data(), W, H) == 0);
}