/*
 * Copyright (c) Camden Dixie O'Brien
 * SPDX-License-Identifier: AGPL-3.0-only
 */

#version 450

#include "config.h"

#define NCOLS 8

layout(location = 0) out vec4 out_colour;

layout(push_constant) uniform Constants {
	vec2 img_size;
#ifdef JULIA
	vec2 julia;
#endif
	vec2 centre;
	float scale;
#ifdef POWRATE
	float zpow;
#endif
} params;

vec3 incol = vec3(0.000000, 0.014444, 0.027321);

vec3 palette[NCOLS] = vec3[](
	vec3(0.000000, 0.049707, 0.107023),
	vec3(0.025187, 0.064803, 0.177888),
	vec3(0.254152, 0.080220, 0.274677),
	vec3(0.502886, 0.080220, 0.278894),
	vec3(1.000000, 0.124772, 0.119538),
	vec3(1.000000, 0.234551, 0.030713),
	vec3(1.000000, 0.381326, 0.000000),
	vec3(1.000000, 0.651406, 0.215861)
);

#if SS == 4
vec2 ss_offsets[SS] = vec2[](
	vec2(-0.25, -0.25),
	vec2( 0.25, -0.25),
	vec2(-0.25,  0.25),
	vec2( 0.25,  0.25)
);
#elif SS == 9
vec2 ss_offsets[SS] = vec2[](
    vec2(-0.33, -0.33), vec2(0.0, -0.33), vec2(0.33, -0.33),
    vec2(-0.33,  0.0),  vec2(0.0,  0.0),  vec2(0.33,  0.0),
    vec2(-0.33,  0.33), vec2(0.0,  0.33), vec2(0.33,  0.33)
);
#else
#error "Unsupported supersampling count"
#endif

vec2 maptoz(vec2 xy)
{
	float aspect = params.img_size.x / params.img_size.y;
	vec2 z = 4.0 * (xy / params.img_size.xy - 0.5);
	z.x *= aspect;
	return params.scale * z + params.centre;
}

vec2 zpow(vec2 z, float p) {
    float r = length(z);
    float theta = atan(z.y, z.x);
    float mag = pow(r, p);
    float phi = p * theta;
    return mag * vec2(cos(phi), sin(phi));
}

int fractal(vec2 c) {
#ifdef JULIA
	vec2 z = c;
	c = params.julia;
#else
	vec2 z = vec2(0.0, 0.0);
#endif
	for (int i = 0; i < MAXITER; ++i) {
#ifdef POWRATE
		z = zpow(z, params.zpow) + c;
#else
		z = vec2(z.x * z.x - z.y * z.y, 2.0 * z.x * z.y) + c;
#endif
		if (length(z) > 2.0)
			return i;
	}
	return -1;
}

vec3 colour(int i) {
	if (i == -1) {
		return incol;
	} else {
		float c = float(NCOLS) * sqrt(float(i) / float(MAXITER));
		int fc = int(floor(c));
		int cc = int(ceil(c));
		float p = c - float(fc);
		p = smoothstep(0.0, 1.0, p);
		return mix(palette[fc], palette[cc], p);
	}
}

void main()
{
	vec2 c;
	vec3 col = vec3(0.0, 0.0, 0.0);
	for (int i = 0; i < SS; ++i) {
		c = maptoz(gl_FragCoord.xy + ss_offsets[i]);
		col += colour(fractal(c));
	}
	col /= SS;
	out_colour = vec4(col, 1.0);
}