import pyglet
mandelbrot_compute_source = """#version 430 core

uniform int u_maxIter;
uniform vec2 u_resolution;
uniform vec2 u_real_range;
uniform vec2 u_imag_range;

layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
layout(location = 0, rgba32f) uniform image2D img_output;

int mandelbrot({vec2type} c, int maxIter) {
    {vec2type} z = {vec2type}(0.0, 0.0);
    for (int n = 0; n < maxIter; n++) {
        if (dot(z, z) > 4.0) {
            return n;
        }
        z = {vec2type}(
            z.x * z.x - z.y * z.y + c.x,
            2.0 * z.x * z.y + c.y
        );
    }
    return maxIter;
}

{vec2type} map_pixel({floattype} x, {floattype} y, {vec2type} resolution, {vec2type} real_range, {vec2type} imag_range) {
    {floattype} real = real_range.x + (x / resolution.x) * (real_range.y - real_range.x);
    {floattype} imag = imag_range.x + (y / resolution.y) * (imag_range.y - imag_range.x);
    return {vec2type}(real, imag);
}

void main() {
    ivec2 texel_coord = ivec2(gl_GlobalInvocationID.xy);
    {vec2type} c = map_pixel({floattype}(texel_coord.x), {floattype}(texel_coord.y), u_resolution, u_real_range, u_imag_range);
    int iters = mandelbrot(c, u_maxIter);

    vec4 value = vec4(0.0, 0.0, 0.0, 1.0);

    if (iters != u_maxIter) {
        float t = float(iters) / float(u_maxIter);
        float pow_amount = 0.7;
        t = pow(t, pow_amount);

        value.r = 9.0 * (1.0 - t) * t * t * t;
        value.g = 15.0 * (1.0 - t) * (1.0 - t) * t * t;
        value.b = 8.5 * (1.0 - t) * (1.0 - t) * (1.0 - t) * t;
    }

    imageStore(img_output, texel_coord, value);
}
"""

def create_mandelbrot_shader(width, height, precision="single"):
    shader_source = mandelbrot_compute_source

    if precision == "single":
        shader_source = shader_source.replace("{vec2type}", "vec2").replace("{floattype}", "float")
    elif precision == "double":
        shader_source = shader_source.replace("{vec2type}", "dvec2").replace("{floattype}", "double")
    else:
        raise TypeError("Invalid Precision")

    shader_program = pyglet.graphics.shader.ComputeShaderProgram(shader_source)

    mandelbrot_image = pyglet.image.Texture.create(width, height, internalformat=pyglet.gl.GL_RGBA32F)

    uniform_location = shader_program['img_output']
    mandelbrot_image.bind_image_texture(unit=uniform_location)

    return shader_program, mandelbrot_image