struct Params {
	max_reflections: i32,
	min_strength: f32,
	sphere_count: i32,
	seed: u32,

	right: vec3f,
	width: f32,
	up: vec3f,
	height: f32,
	forward: vec3f,
	aperture: f32,
	eye: vec3f,
	antifocal: f32,
}

struct Sphere {
	center: vec3f,
	radius: f32,
	emit_color: vec3f,
	reflect_color: vec3f,
	glossiness: f32,
}

struct Vertex {
	@location(0) screen: vec2f,
}

struct Varying {
	@location(0) dir: vec3f,
	@builtin(position) screen: vec4f,
}

var<push_constant> params: Params;
@group(0) @binding(1) var<storage, read> spheres: array<Sphere>;
@group(1) @binding(0) var env_sampler: sampler;
@group(1) @binding(1) var env_texture: texture_cube<f32>;

@vertex
fn on_vertex(in: Vertex) -> Varying {
	let m = mat3x3(params.width * params.right, params.height * params.up, params.forward);
	return Varying(m * vec3(in.screen, 1.0), vec4(in.screen, 0.0, 1.0));
}

@fragment
fn on_fragment(in: Varying) -> @location(0) vec4f {
	return vec4(trace_fragment(in), 1.);
}

fn sqr(v: vec3f) -> f32 {
        return dot(v, v);
}

var<private> pos: vec3f;
var<private> ray: vec3f;

fn to_sphere(center: vec3f, radius: f32, t: ptr<function, f32>) -> bool {
	let c = sqr(pos - center) - radius * radius;
	let b = 2 * dot(pos - center, ray);
	let a = sqr(ray);

	let D = b * b - 4 * a * c;
	if (D <= 0) {
		return false;
	}
	*t = (- b - sqrt(D)) / (2 * a);
	if (*t < 0) {
		return false;
	}
	return true;
}

fn trace_fragment(in: Varying) -> vec3f {
	seed(in.screen);

	var result = vec3(0.);
	var color = vec3(1.);
	let view_mtx = mat3x3(params.right, params.up, params.forward);
	let aperture_offset_rel = params.aperture * rand_disc();
	let aperture_offset_abs = view_mtx * vec3(aperture_offset_rel, 0.);
	pos = params.eye + aperture_offset_abs;
	let off_px = vec2(rand_float(), rand_float()) - .5;
	let off_w = mat2x3(dpdx(in.dir), dpdy(in.dir));
	let dir = in.dir + off_w * off_px;
	ray = normalize(dir - params.antifocal * aperture_offset_abs);

	for (var k = 0; k < params.max_reflections; k++) {
		var sphere = -1;
		var t = 1.0e9;
		for (var k = 0; k < params.sphere_count; k++) {
			var t1: f32;
			if (to_sphere(spheres[k].center, spheres[k].radius, &t1) && t1 < t) {
					sphere = k;
					t = t1;
			}
		}
		if (sphere == -1) {
			let env = textureSampleLevel(env_texture, env_sampler, ray, 0.0);
			result += 3.0 * color * env.xyz;
			break;
		}
		let s = spheres[sphere];
		pos += t * ray;
		let normal = (pos - s.center) / s.radius;
		if (all(s.emit_color == vec3(0.))) {
			color *= s.reflect_color;
			let diffuse = normal + rand_sphere();
			let specular = reflect(ray, normal);
			ray = normalize(mix(diffuse, specular, s.glossiness));
		} else {
			let d = dot(-ray, normal);
			let strength = d * d; // it would be 1 for surface emission, but this models volume emission
			result += color * s.emit_color * strength;
			color *= (1. - strength);
			pos += 1e-3 * ray;
		}
		if (length(color) < params.min_strength) {
			break;
		}
	}

	return result;
}

fn hash(key : u32) -> u32 {
        var v = key;
        v *= 0xb384af1bu;
        v ^= v >> 15u;
        return v;
}

var<private> rand_state: u32;

fn seed(key: vec4f) {
	let x = bitcast<u32>(key.x);
	let y = bitcast<u32>(key.y);
	rand_state = hash(hash(hash(params.seed) ^ x) ^ y);
}

fn rand_next() -> u32 {
	rand_state = hash(rand_state);
	return rand_state;
}

fn rand_float() -> f32 {
	return f32(rand_next()) / 0x1p32;
}

fn rand_disc() -> vec2f {
	for (var k = 0; k < 16; k++) {
		let v = vec2f(rand_float(), rand_float()) - 0.5;
		if (length(v) <= 0.5) {
			return 2. * v;
		}
	}
	return vec2f(0.0); // safeguard
}

fn rand_sphere() -> vec3f {
	for (var k = 0; k < 16; k++) {
		let v = vec3f(rand_float(), rand_float(), rand_float()) - 0.5;
		let l = length(v);
		if (length(v) <= 0.5) {
			return v / l;
		}
	}
	return vec3f(0.0); // safeguard
}