#version 450
/* Copyright (c) 2020-2021, Holochip
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 the "License";
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

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

layout(binding = 0, rg8ui) uniform readonly uimage2D input_frequency;
layout(binding = 1, r8ui) uniform writeonly uimage2D output_sampling_rate;

layout(binding = 2) buffer FrequencyInformation
{
	uvec4 settings;                     // {frame_width, frame_height, output_width, output_height}
	uvec4 frequency_information;        // { max_rate_x, max_rate_y, n_rates, ...}
	uvec2 rates[];
}
params;

void main()
{
	const uint x0            = gl_GlobalInvocationID.x,
	           y0            = gl_GlobalInvocationID.y,
	           frame_width   = params.settings[0],
	           frame_height  = params.settings[1],
	           output_width  = params.settings[2],
	           output_height = params.settings[3];

	const uint delta_x = max(1, uint(round(float(frame_width) / float(output_width)))),
	           delta_y = max(1, uint(round(float(frame_height) / float(output_height))));
	if (x0 >= output_width || y0 >= output_height)
	{
		return;
	}

	vec2 max_freqs = vec2(0, 0);
	for (uint i = 0; i < delta_x; ++i)
	{
		for (uint j = 0; j < delta_y; ++j)
		{
			ivec2 coord = ivec2(delta_x * x0 + i, delta_y * y0 + j);
			vec2  freq  = vec2(imageLoad(input_frequency, coord)) / 255.f;
			max_freqs   = max(max_freqs, freq);
		}
	}
	
	vec2 freqs = min(1.25 * sqrt(max_freqs), vec2(1, 1));

	const float min_rate = 1,
	            max_rate = max(params.frequency_information.x, params.frequency_information.y);

	const vec2 optimal_rate = freqs * vec2(min_rate, min_rate) + (1 - freqs) * vec2(max_rate, max_rate);        //round(frequency * min_rate + (1 - frequency) * max_rate);

	const uint n_rates            = params.frequency_information.z;
	uint       optimal_rate_index = 0;
	float      current_cost       = 1 + 2 * max_rate * max_rate;
	for (uint i = 0; i < n_rates; ++i)
	{
		uvec2 rate = params.rates[i];
		float cost = (rate.x - optimal_rate.x) * (rate.x - optimal_rate.x) + (rate.y - optimal_rate.y) * (rate.y - optimal_rate.y);
		if (cost < current_cost)
		{
			current_cost       = cost;
			optimal_rate_index = i;
		}
	}

	uint optimal_rate_x = params.rates[optimal_rate_index].x, optimal_rate_y = params.rates[optimal_rate_index].y;
	uint rate_code = uint(optimal_rate_y >> 1) | ((optimal_rate_x << 1) & 12);
	imageStore(output_sampling_rate, ivec2(x0, y0), uvec4(rate_code));
}