// Copyright 2021-2022, Collabora Ltd.
// Author: Jakob Bornecrantz <jakob@collabora.com>
// SPDX-License-Identifier: BSL-1.0

#version 460
#extension GL_GOOGLE_include_directive : require

#include "srgb.inc.glsl"


// The size of the distortion texture dimensions in texels.
layout(constant_id = 0) const int distortion_texel_count = 2;

// Should we do timewarp.
layout(constant_id = 1) const bool do_timewarp = false;

layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

layout(set = 0, binding = 0) uniform sampler2D source[2];
layout(set = 0, binding = 1) uniform sampler2D distortion[6];
layout(set = 0, binding = 2) uniform writeonly restrict image2D target;
layout(set = 0, binding = 3, std140) uniform restrict Config
{
	ivec4 views[2];
	vec4 pre_transform[2];
	vec4 post_transform[2];
	mat4 transform_timewarp_scanout_begin[2];
	mat4 transform_timewarp_scanout_end[2];
}ubo;


vec2 position_to_uv(ivec2 extent, uint ix, uint iy)
{
	// Turn the index into floating point.
	vec2 xy = vec2(float(ix), float(iy));

	// The inverse of the extent of the target image is the pixel size in [0 .. 1] space.
	vec2 extent_pixel_size = vec2(1.0 / float(extent.x), 1.0 / float(extent.y));

	// Per-target pixel we move the size of the pixels.
	vec2 dist_uv = xy * extent_pixel_size;

	// Emulate a triangle sample position by offset half target pixel size.
	dist_uv = dist_uv + extent_pixel_size / 2.0;


	// To correctly sample we need to put position (0, 0) in the
	// middle of the (0, 0) texel in the distortion textures. That's why we
	// offset with half the texel size, pushing all samples into the middle
	// of each texels, a kin to a vertex buffer. We need to put uv coord
	// (1, 1) in the middle of the last texel, that pixel is (size - 1)
	// texels away from the first texel. So we need to scale [0 .. 1] to
	// [0 .. size - 1].

#define DIM (float(distortion_texel_count))
#define STRETCH ((DIM - 1.0) / DIM)
#define OFFSET (1.0 / (DIM * 2.0))

	dist_uv = (dist_uv * STRETCH) + OFFSET;

	return dist_uv;
}

vec2 transform_uv_subimage(vec2 uv, uint iz)
{
	vec2 values = uv;

	// To deal with OpenGL flip and sub image view.
	values.xy = values.xy * ubo.post_transform[iz].zw + ubo.post_transform[iz].xy;

	// Ready to be used.
	return values.xy;
}

vec2 transform_uv_timewarp(vec2 uv, uint iz)
{
	vec4 values = vec4(uv, -1, 1);

	// From uv to tan angle (tangent space).
	values.xy = values.xy * ubo.pre_transform[iz].zw + ubo.pre_transform[iz].xy;
	values.y = -values.y; // Flip to OpenXR coordinate system.

	// Timewarp including scanline timewarp for rolling refresh panels.
	values = ubo.transform_timewarp_scanout_begin[iz] * values * (1 - uv.y) +
	         ubo.transform_timewarp_scanout_end[iz] * values * uv.y;
	values.xy = values.xy * (1.0 / max(values.w, 0.00001));

	// From [-1, 1] to [0, 1]
	values.xy = values.xy * 0.5 + 0.5;

	// Done.
	return values.xy;
}

vec2 transform_uv(vec2 uv, uint iz)
{
	if (do_timewarp) {
		uv = transform_uv_timewarp(uv, iz);
	}

	return transform_uv_subimage(uv, iz);
}

void main()
{
	uint ix = gl_GlobalInvocationID.x;
	uint iy = gl_GlobalInvocationID.y;
	uint iz = gl_GlobalInvocationID.z;

	ivec2 offset = ivec2(ubo.views[iz].xy);
	ivec2 extent = ivec2(ubo.views[iz].zw);

	if (ix >= extent.x || iy >= extent.y) {
		return;
	}

	vec2 dist_uv = position_to_uv(extent, ix, iy);

	vec2 r_uv = texture(distortion[iz + 0], dist_uv).xy;
	vec2 g_uv = texture(distortion[iz + 2], dist_uv).xy;
	vec2 b_uv = texture(distortion[iz + 4], dist_uv).xy;

	// Do any transformation needed.
	r_uv = transform_uv(r_uv, iz);
	g_uv = transform_uv(g_uv, iz);
	b_uv = transform_uv(b_uv, iz);

	// Sample the source with distorted and chromatic-aberration corrected samples.
	vec4 colour = vec4(              //
	    texture(source[iz], r_uv).r, //
	    texture(source[iz], g_uv).g, //
	    texture(source[iz], b_uv).b, //
	    1);                          //

	// Do colour correction here since there are no automatic conversion in hardware available.
	colour = vec4(from_linear_to_srgb(colour.rgb), 1);

	imageStore(target, ivec2(offset.x + ix, offset.y + iy), colour);
}