#include "Image.h"

#include <fstream>
#include <iostream>

#include "../tools/Threadpool.h"
#include "StratifiedSampler.h"

namespace util {
Image::Image(int width, int height) : width(width), height(height) {
	Vec3 color({});
	for (int i = 0; i < width * height; i++) {
		vec.insert(vec.end(), color);
	}
}
/*
void Image::setPixels(std::shared_ptr<Sampler> sampler) {
    Threadpool tp(4);
    for (int x = 0; x != width; x++) {
        for (int y = 0; y != height; y++) {
            vec[width * y + x] = sampler->color(x, y);
        }
    }
} */
void Image::setPixel(int x, int y, Vec3 color) {
	vec[width * y + x] = color;
}
void Image::setPixels(size_t threadcount, std::shared_ptr<Sampler> sampler) {
	Threadpool tp(threadcount);
	for (int x = 0; x != width; x++) {
		for (int y = 0; y != height; y++) {
			tp.queueTask(std::bind([this, x, y, sampler]() {
				this->setPixelsTask(x, y, sampler);
			}));
		}
	}
	// std::cout << "Done queing" << std::endl;
}

void Image::setPixelsTask(int x, int y, std::shared_ptr<Sampler> sampler) {
	Vec3 v = sampler->color(x, y);
	setPixel(x, y, v);
}

void Image::gammaCorrect(float gamma) {
	// correct the whole data-array with the given gamma
	std::transform(vec.begin(), vec.end(), vec.begin(),
	               [gamma](util::Vec3 v) -> util::Vec3 {
		               return util::Vec3(std::powf(v.x(), 1 / gamma),
		                                 std::powf(v.y(), 1 / gamma),
		                                 std::powf(v.z(), 1 / gamma));
	               });
}
Vec3 Image::operator[](const std::array<int, 2>& i) const {
	return vec[width * i[1] + i[0]];
}
Vec3& Image::operator[](const std::array<int, 2>& i) {
	return vec[width * i[1] + i[0]];
}
Image raytrace(size_t threadcount, const cam::CamObs& cam,
               const std::shared_ptr<Sampler>& sampler, size_t n) {
	Image result(cam.width, cam.height);

	result.setPixels(threadcount, std::make_shared<StratifiedSampler>(
	                                  StratifiedSampler(sampler, n))
	                 // sampler
	);
	result.gammaCorrect(2.2);
	return result;
}
void writeBmp(const char* filename, Image img) {
	int width = img.width;    // Width of the image
	int height = img.height;  // Height of the image

	int horizontalBytes;  // Horizontal bytes to add, so the pixelarrays width
	                      // is a multiple of 4
	horizontalBytes = 4 - ((width * 3) % 4);
	if (horizontalBytes == 4) horizontalBytes = 0;

	int pixelArraySize;  // The size of the pixelArray
	pixelArraySize = ((width * 3) + horizontalBytes) * height;

	// Headers
	int header[12];
	// Bitmap file header
	char bb = 'B';  // bfType
	char mm = 'M';
	header[0] = pixelArraySize + 54;  // bfSize
	header[1] = 0;                    // bfReserved
	header[2] = 54;                   // bfOffbits
	// Bitmap information header
	header[3] = 40;              // biSize
	header[4] = width;           // biWidth
	header[5] = height;          // biHeight
	short biPLanes = 1;          // biPlanes
	short biBitCount = 24;       // biBitCount
	header[6] = 0;               // biCompression
	header[7] = pixelArraySize;  // biSizeImage
	header[8] = 0;               // biXPelsPerMeter
	header[9] = 0;               // biYPelsPerMeter
	header[10] = 0;              // biClrUsed
	header[11] = 0;              // biClrImportant

	std::ofstream ofile(filename, std::ios::binary);

	// Write the header in the right order
	// bfType, ...
	ofile.write(&bb, sizeof(bb));
	ofile.write(&mm, sizeof(mm));
	// ... bfSize, BfReserved, bfOffbits, biSize, biWidth, bitHeight, ...
	for (int i = 0; i < 6; i++) {
		ofile.write((char*)&header[i], sizeof(header[i]));
	}
	// ... biPlanes, bitBitCount, ...
	ofile.write((char*)&biPLanes, sizeof(biPLanes));
	ofile.write((char*)&biBitCount, sizeof(biBitCount));
	// ... biCompression, biSizeImage, biXPelsPerMeter, biYPelsPerMeter,
	// biClrUsed, biClrImportant
	for (int i = 6; i < 12; i++) {
		ofile.write((char*)&header[i], sizeof(header[i]));
	}
	// The colors can only have a value from 0 to 255, so a char is enough to
	// store them
	char blue, green, red;
	// Bmp is written from top to bottom
	for (int y = height - 1; y >= 0; y--) {
		for (int x = 0; x <= width - 1; x++) {
			red = std::clamp<float>((img[{x, y}][0]), 0, 1) * 255;
			green = std::clamp<float>((img[{x, y}][1]), 0, 1) * 255;
			blue = std::clamp<float>((img[{x, y}][2]), 0, 1) * 255;

			// bmp colors are bgr not rgb
			char bgr[3] = {blue, green, red};
			for (int i = 0; i < 3; i++) {
				char c0 = (bgr[i] & 0x00FF);
				// char c8 = ((bgr[i] & (unsigned int)0xFF00) >> 8);

				ofile.write(&c0, sizeof(c0));
				// ofile.write (&c8, sizeof (c8));
			}
		}
		// If needed add extra bytes after each row
		if (horizontalBytes != 0) {
			char null;
			for (int n = 0; n < horizontalBytes; n++) {
				ofile.write(&null, sizeof(null));
			}
		}
	}

	ofile.close();
}
}  // namespace util