#include <cstdint>
#include <cstdio>
#include <cstdlib>

struct Pixel {
  uint8_t r, g, b;
};

struct Complex {
  float r, i;

  __device__ float norm2() const { return r * r + i * i; }
};

__device__ Complex operator*(Complex a, Complex b) {
  return Complex{(a.r * b.r) - (a.i * b.i), (a.i * b.r) + (a.r * b.i)};
}

__device__ Complex operator+(Complex a, Complex b) {
  return Complex{a.r + b.r, a.i + b.i};
}

__device__ int julia(int width, int height, int x, int y) {
  constexpr float scale = 1.5;
  constexpr int   N     = 200;

  const float jx = scale * (width / 2 - x) / (width / 2);
  const float jy = scale * (height / 2 - y) / (height / 2);

  const Complex c{-0.8, 0.156};
  Complex       a{jx, jy};

  for (int i = 0; i < N; ++i) {
    a = a * a + c;
    if (a.norm2() > 1000) {
      return 0;
    }
  }
  return 1;
}

__global__ void juliaMain(int width, int height, Pixel* image) {
  const int x = blockIdx.x;
  const int y = blockIdx.y;

  constexpr Pixel background{41, 95, 152};
  constexpr Pixel juliaColour{228, 192, 135};

  const Pixel pixel =
      julia(width, height, x, y) == 1 ? juliaColour : background;

  image[y * width + x] = pixel;
}

bool write_pbm(const Pixel* image, int width, int height, const char* path) {
  const size_t num_pixels = width * height;

  FILE* file = fopen(path, "wb");
  if (!file) {
    return false;
  }

  fprintf(file, "P6\n%d %d\n255\n", width, height);
  if (fwrite(image, sizeof(Pixel), num_pixels, file) != num_pixels) {
    fclose(file);
    return false;
  }

  fclose(file);
  return true;
}

int main(int argc, const char** argv) {
  const int width  = argc > 1 ? atoi(argv[1]) : 1920;
  const int height = argc > 2 ? atoi(argv[2]) : 1080;

  bool success = false;

  const dim3 dim(width, height);
  const int  image_size_bytes = width * height * sizeof(Pixel);
  auto       image_host       = new Pixel[width * height];
  Pixel*     image_dev        = nullptr;

  if (cudaMalloc(&image_dev, image_size_bytes) != cudaSuccess) {
    goto cleanup;
  }

  juliaMain<<<dim, 1>>>(width, height, image_dev);

  if (cudaMemcpy(
          image_host, image_dev, image_size_bytes, cudaMemcpyDeviceToHost) !=
      cudaSuccess) {
    goto cleanup;
  }

  if (!write_pbm(image_host, width, height, "julia.pbm")) {
    goto cleanup;
  }

  success = true;

cleanup:
  delete[] image_host;
  if (image_dev) {
    cudaFree(image_dev);
  }
  return success ? 0 : 1;
}