/* main.c
 * Created: 2023-09-02
 * Author: Aryadev Chavali
 * Description: Entrypoint of program
 */

#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <pthread.h>
#include <unistd.h>

#include <raylib.h>
#include <raymath.h>

#include <stb/stb_image_write.h>

#define WIDTH       1024
#define HEIGHT      1024
#define MAX_THREADS 4
uint64_t MAX_ITER = 1 << 2;

Color cells[WIDTH * HEIGHT];
pthread_t threads[MAX_THREADS];

#define SQUARE(a) (a * a)

Color iter_to_colour(size_t iterations)
{
  if (iterations == MAX_ITER)
    return BLACK;

  // quadratic ratio
  float ratio = Remap(iterations * (MAX_ITER - iterations), 0,
                      SQUARE(MAX_ITER) / 4, 0, 1);

  return (Color){255 * SQUARE(ratio), 255 * SQUARE(SQUARE(ratio)), 255 * ratio,
                 255};
}

struct ThreadArg
{
  size_t x_0, x_t;
  size_t y_0, y_t;
  bool done;
};

void *generate_colours(void *state)
{
  struct ThreadArg *ptr = (struct ThreadArg *)state;
  ptr->done             = false;
  for (size_t i = ptr->x_0; i < ptr->x_t; ++i)
    for (size_t j = ptr->y_0; j < ptr->y_t; ++j)
    {
      Vector2 init   = {Remap(i, 0, WIDTH, -1.5, 0.5),
                        Remap(j, 0, HEIGHT, -1, 1)};
      Vector2 update = {0};
      size_t iterations;
      for (iterations = 0;
           iterations < MAX_ITER && Vector2LengthSqr(update) <= 4; ++iterations)
      {
        double new_x = SQUARE(update.x) - SQUARE(update.y) + init.x;
        update.y     = (2 * update.x * update.y) + init.y;
        update.x     = new_x;
      }
      cells[(i * WIDTH) + j] = iter_to_colour(iterations);
    }
  ptr->done = true;
  return NULL;
}

void threads_start_render(struct ThreadArg *args)
{
  for (size_t i = 0; i < MAX_THREADS; ++i)
    pthread_create(&threads[i], NULL, &generate_colours, &args[i]);
}

bool threads_done(struct ThreadArg *args)
{
  for (size_t i = 0; i < MAX_THREADS; ++i)
    if (!args[i].done)
      return false;
  return true;
}

void threads_cancel_render(void)
{
  for (size_t i = 0; i < MAX_THREADS; ++i)
    pthread_join(threads[i], NULL);
}

bool write_to_png(const char *filepath)
{
  unsigned char *image_data = calloc(3 * WIDTH * HEIGHT, sizeof(*image_data));

  size_t image_ptr = 0;

  for (size_t i = 0; i < HEIGHT; ++i)
    for (size_t j = 0; j < WIDTH; ++j, image_ptr += 3)
    {
      Color c                 = cells[(j * WIDTH) + i];
      unsigned char colour[3] = {c.r, c.g, c.b};
      memcpy(image_data + image_ptr, colour, sizeof(*colour) * 3);
    }

  stbi_write_png(filepath, WIDTH, HEIGHT, 3, image_data, 3 * WIDTH);
  free(image_data);
  return true;
}

#define ZOOM_INC 0.3f
#define MOVE_INC WIDTH / 50
int main(void)
{
  InitWindow(WIDTH, HEIGHT, "Mandelbrot simulation");
  Camera2D camera = {0};
  camera.zoom     = 1.0f;

  Image img         = {.data    = cells,
                       .width   = WIDTH,
                       .height  = HEIGHT,
                       .format  = PIXELFORMAT_UNCOMPRESSED_R8G8B8A8,
                       .mipmaps = 1};
  Texture2D texture = LoadTextureFromImage(img);

  SetTargetFPS(60);

  struct ThreadArg args[] = {{0, WIDTH / 2, 0, HEIGHT, false},
                             {WIDTH / 2, WIDTH, 0, HEIGHT / 4, false},
                             {WIDTH / 2, WIDTH, HEIGHT / 4, HEIGHT / 2, false},
                             {WIDTH / 2, WIDTH, HEIGHT / 2, HEIGHT, false}};
  threads_start_render(args);

  const size_t delta = 10;
  bool rendered      = false;

  for (size_t prev = 0, ticks = 0; !WindowShouldClose(); ++ticks)
  {
    if (IsKeyPressed(KEY_UP) || IsKeyDown(KEY_UP))
      camera.target.y -= MOVE_INC;
    else if (IsKeyPressed(KEY_DOWN) || IsKeyDown(KEY_DOWN))
      camera.target.y += MOVE_INC;
    else if (IsKeyPressed(KEY_RIGHT) || IsKeyDown(KEY_RIGHT))
      camera.target.x += MOVE_INC;
    else if (IsKeyPressed(KEY_LEFT) || IsKeyDown(KEY_LEFT))
      camera.target.x -= MOVE_INC;
    else if (IsKeyPressed(KEY_N) || IsKeyDown(KEY_N))
    {
      camera.zoom += ZOOM_INC;
      if (camera.zoom > 3.0f)
        camera.zoom = 3.0f;
    }
    else if (IsKeyPressed(KEY_M) || IsKeyDown(KEY_M))
    {
      camera.zoom -= ZOOM_INC;
      if (camera.zoom < 0.1f)
        camera.zoom = 0.1f;
    }
    else if (IsKeyPressed(KEY_SPACE))
    {
      threads_cancel_render();
      memset(cells, 0, WIDTH * HEIGHT);
      MAX_ITER *= 2;
      rendered = false;
      threads_start_render(args);
    }

    if (ticks - prev > delta)
    {
      prev      = ticks;
      bool done = threads_done(args);
      if (!done || !rendered)
        UpdateTexture(texture, cells);
      else if (done)
        rendered = true;
    }

    char iters[128];
    sprintf(iters, "iterations=%lu", MAX_ITER);

    BeginDrawing();
    BeginMode2D(camera);
    ClearBackground(BLACK);
    DrawTexturePro(texture, (Rectangle){0, 0, WIDTH, HEIGHT},
                   (Rectangle){WIDTH / 2, HEIGHT / 2, WIDTH, HEIGHT},
                   (Vector2){WIDTH / 2, HEIGHT / 2}, -90, WHITE);
    DrawText(iters, 200, 200, 50, RAYWHITE);
    EndMode2D();
    EndDrawing();
  }
  write_to_png("test.png");
  threads_cancel_render();
  UnloadTexture(texture);
  CloseWindow();
  return 0;
}