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

#version 450

#include "config.h"

#define NCOLS 8

#if !defined(CPOW_DTHETA) && !defined(THING)
#define MAXMAG 2.0
#endif

layout(location = 0) out vec4 out_colour;

layout(push_constant) uniform Constants {
	vec2 img_size;
#if defined(JULIA) || defined(THING)
	vec2 julia;
#endif
	vec2 centre;
#ifdef CPOW_DTHETA
	vec2 cpow;
#endif
#ifdef POWRATE
	float zpow;
#endif
	float scale;
} 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 cmul(vec2 a, vec2 b)
{
	return vec2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
}

vec2 cexp(vec2 z)
{
	return exp(z.x) * vec2(cos(z.y), sin(z.y));
}

vec2 cln(vec2 z)
{
	return vec2(log(length(z)), atan(z.y, z.x));
}

vec2 cpow(vec2 z, vec2 p)
{
	if (z.x == 0.0 && z.y == 0.0) return vec2(0.0, 0.0);
	return cexp(cmul(p, cln(z)));
}

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

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 CPOW_DTHETA
		z = cpow(z, params.cpow) + c;
#elif defined(POWRATE)
		z = zpow(z, params.zpow) + c;
#elif defined(THING)
		z = cpow(z, c) + params.julia;
#else
		z = vec2(z.x * z.x - z.y * z.y, 2.0 * z.x * z.y) + c;
#endif
		if (length(z) > MAXMAG)
			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);
}