from pathlib import Path
from time import time

# Generates one binary PPM image of the Mandelbrot set.
# The image size is calibrated so the Python implementation stays around 1-2 seconds
# on a typical development machine while still doing enough numeric work.

WIDTH = 1000
HEIGHT = 700
MAX_ITER = 96


def benchmark_folder() -> Path:
    folder = Path.cwd() / "samples" / "perf"
    if folder.is_dir():
        return folder
    return Path.cwd()


def main() -> None:
    width = WIDTH
    height = HEIGHT
    max_iter = MAX_ITER

    begin = time()

    pixels = bytearray(width * height * 3)
    index = 0
    checksum = 0
    scale_x = 3.3 / (width - 1)
    scale_y = 2.4 / (height - 1)

    for y in range(height):
        ci = -1.2 + y * scale_y
        for x in range(width):
            cr = -2.1 + x * scale_x
            zr = 0.0
            zi = 0.0
            iteration = 0

            while zr * zr + zi * zi <= 4.0 and iteration < max_iter:
                next_zr = zr * zr - zi * zi + cr
                zi = 2.0 * zr * zi + ci
                zr = next_zr
                iteration += 1

            if iteration == max_iter:
                r = 0
                g = 0
                b = 0
            else:
                shade = iteration * 255 // max_iter
                r = shade
                g = shade * 7 // 10
                b = 255 - shade

            pixels[index] = r
            pixels[index + 1] = g
            pixels[index + 2] = b
            checksum += r + g + b
            index += 3

    output = benchmark_folder() / "Mandelbrot.python.ppm"
    with output.open("wb") as stream:
        stream.write(f"P6\n{width} {height}\n255\n".encode("ascii"))
        stream.write(pixels)

    end = time()
    print(f"image = {output}")
    print(f"checksum = {checksum}")
    print(f"Duration = {int((end - begin) * 1000)} ms")


if __name__ == "__main__":
    main()
