#include "camera.h"

#include "ray.h"
#include "rng.h"

#include <float.h>
#include <math.h>
#include <threads.h>

#define MAX_ITER 10
#define MIN_T 1e-6

#define NSAMPLES 100
#define SAMPLE_WEIGHT (1.0 / (double)NSAMPLES)
#define SAMPLE_STDDEV 0.333

#define GAMMA 2.2

#define NTHREADS 40

typedef struct {
	const camera_t *camera;
	const obj_t *scene;
	pix_t *pixels;
	rng_t rng;
	unsigned start_y, row_count, scene_count;
} work_slice_t;

static const vec3_t lightblue = { 0.4, 0.6, 1.0 };
static const vec3_t white = { 1.0, 1.0, 1.0 };
static const vec3_t black = { 0.0, 0.0, 0.0 };

static ray_t scatter(hit_t hit, rng_t *rng)
{
	return (ray_t) {
		.orig = hit.point,
		.dir = vec3_add(hit.normal, rng_vec3(rng)),
	};
}

static vec3_t
trace(ray_t ray, const obj_t *scene, unsigned scene_count, rng_t *rng)
{
	double coeff = 1.0;
	for (unsigned i = 0; i < MAX_ITER; ++i) {
		hit_t hit = { .t = DBL_MAX };
		for (unsigned j = 0; j < scene_count; ++j)
			scene[j].intersect(scene[j].params, ray, &hit, MIN_T, hit.t);

		if (hit.t == DBL_MAX) {
			const double a = (ray.dir.y + 1.0) / 2.0;
			return vec3_scale(
			    vec3_add(vec3_scale(lightblue, a), vec3_scale(white, 1 - a)),
			    coeff);
		}

		ray = scatter(hit, rng);
		coeff *= 0.3;
	}

	return black;
}

static double linear_to_gamma(double channel)
{
	return pow(fmin(channel, 1.0), 1.0 / GAMMA);
}

static void setpix(vec3_t col, pix_t *out)
{
	out->r = UINT16_MAX * linear_to_gamma(col.x);
	out->g = UINT16_MAX * linear_to_gamma(col.y);
	out->b = UINT16_MAX * linear_to_gamma(col.z);
	out->a = UINT16_MAX;
}

static int render_thread(void *arg)
{
	work_slice_t *slice = (work_slice_t *)arg;
	const camera_t *camera = slice->camera;
	const uint32_t w = camera->img_width;

	const unsigned stop_y = slice->start_y + slice->row_count;
	for (unsigned y = slice->start_y; y < stop_y; ++y) {
		const vec3_t row
		    = vec3_add(camera->pix_origin, vec3_scale(camera->y_step, y));
		for (unsigned x = 0; x < w; ++x) {
			const vec3_t pix = vec3_add(row, vec3_scale(camera->x_step, x));

			vec3_t colour = black;
			for (unsigned i = 0; i < NSAMPLES; ++i) {
				const vec3_t jitter
				    = rng_gaussian_xy(&slice->rng, SAMPLE_STDDEV);
				const vec3_t offset = vec3_add(
				    vec3_scale(camera->x_step, jitter.x),
				    vec3_scale(camera->y_step, jitter.y));

				const ray_t ray = {
					.orig = camera->pos,
					.dir = vec3_unit(vec3_add(pix, offset)),
				};
				const vec3_t sample = trace(
				    ray, slice->scene, slice->scene_count, &slice->rng);

				colour = vec3_add(colour, vec3_scale(sample, SAMPLE_WEIGHT));
			}

			setpix(colour, slice->pixels + (w * y + x));
		}
	}

	return 0;
}

camera_t camera_init(
    vec3_t pos, double focal_len, double viewport_height, uint32_t img_width,
    uint32_t img_height)
{
	const double aspect = (double)img_width / (double)img_height;
	const double viewport_width = viewport_height * aspect;
	const vec3_t viewport_disp = { 0, 0, focal_len };
	const vec3_t u = { viewport_width, 0, 0 };
	const vec3_t v = { 0, -viewport_height, 0 };
	const vec3_t topleft = vec3_sub(
	    vec3_sub(pos, viewport_disp), vec3_scale(vec3_add(u, v), 0.5));

	const vec3_t x_step = vec3_scale(u, 1.0 / (double)img_width);
	const vec3_t y_step = vec3_scale(v, 1.0 / (double)img_height);
	const vec3_t pix_origin
	    = vec3_add(topleft, vec3_scale(vec3_add(x_step, y_step), 0.5));

	return (camera_t) {
		.pos = pos,
		.pix_origin = pix_origin,
		.x_step = x_step,
		.y_step = y_step,
		.img_width = img_width,
		.img_height = img_height,
	};
}

void camera_render(
    const camera_t *camera, const obj_t *scene, unsigned scene_count,
    img_t *img_out)
{
	img_out->w = camera->img_width;
	img_out->h = camera->img_height;

	const unsigned rows_per_thread = img_out->h / NTHREADS;
	const unsigned rem_rows = img_out->h % NTHREADS;

	thrd_t threads[NTHREADS];
	work_slice_t slices[NTHREADS];
	for (unsigned i = 0; i < NTHREADS; ++i) {
		slices[i].camera = camera;
		slices[i].scene = scene;
		slices[i].scene_count = scene_count;
		slices[i].pixels = img_out->pix;
		slices[i].rng = rng_init(i);

		slices[i].start_y = i * rows_per_thread;
		slices[i].row_count = rows_per_thread;
		if (rem_rows != 0 && i == NTHREADS - 1)
			slices[i].row_count += rem_rows;

		thrd_create(threads + i, render_thread, slices + i);
	}

	for (unsigned i = 0; i < NTHREADS; ++i)
		thrd_join(threads[i], 0);
}