#include "dimension.h"
#include <math.h>

typedef struct physics_sphere {
  dmnsn_vector center;
  dmnsn_vector velocity;
  double radius;
  dmnsn_color color;
} physics_sphere;

physics_sphere
make_sphere(size_t x, size_t y, size_t z, size_t size)
{
  --size;

  double dx = sin(2*M_PI*x/size);
  double dy = sin(2*M_PI*y/size);
  double dz = sin(2*M_PI*z/size);
  dmnsn_vector c = dmnsn_vector_sub(
    dmnsn_vector_add(
      dmnsn_vector_mul(
        5.0/size,
        dmnsn_new_vector(x, y, z)
      ),
      dmnsn_vector_div(
        dmnsn_new_vector(dy + dz, dx + dz, dx + dy),
        4.0
      )
    ),
    dmnsn_new_vector(2.5, 2.5, 2.5)
  );

  double r = 2.0/size;

  dmnsn_color color = dmnsn_color_from_sRGB(
    dmnsn_new_color((double)x/size, (double)y/size, (double)z/size)
  );

  physics_sphere s = {
    .center = c,
    .velocity = dmnsn_zero,
    .radius = r,
    .color = color,
  };
  return s;
}

dmnsn_array *
make_spheres()
{
  const size_t size = 10;
  dmnsn_array *spheres = dmnsn_new_array(sizeof(physics_sphere));
  for (size_t x = 0; x < size; ++x) {
    for (size_t y = 0; y < size; ++y) {
      for (size_t z = 0; z < size; ++z) {
        physics_sphere sphere = make_sphere(x, y, z, size);
        dmnsn_array_push(spheres, &sphere);
      }
    }
  }

  return spheres;
}

dmnsn_scene *
make_scene(dmnsn_array *spheres, dmnsn_canvas *canvas, dmnsn_camera *camera)
{
  dmnsn_scene *scene = dmnsn_new_scene();
  DMNSN_INCREF(canvas);
  scene->canvas = canvas;
  DMNSN_INCREF(camera);
  scene->camera = camera;

  scene->default_texture->pigment = dmnsn_new_solid_pigment(dmnsn_black);
  scene->default_texture->finish.ambient = dmnsn_new_basic_ambient(
    dmnsn_color_from_sRGB(dmnsn_color_mul(0.25, dmnsn_white))
  );
  scene->default_texture->finish.diffuse = dmnsn_new_lambertian(
    dmnsn_sRGB_inverse_gamma(0.8)
  );

  scene->background = dmnsn_new_solid_pigment(
    dmnsn_color_from_sRGB(
      dmnsn_color_mul(0.5, dmnsn_new_color(0.73, 0.90, 0.97))
    )
  );

  DMNSN_ARRAY_FOREACH (physics_sphere *, s, spheres) {
    dmnsn_object *sphere = dmnsn_new_sphere();

    sphere->texture = dmnsn_new_texture();
    sphere->texture->pigment = dmnsn_new_solid_pigment(s->color);
    double maxcomponent = dmnsn_max(
      dmnsn_max(s->color.R, s->color.G),
      s->color.B
    );
    dmnsn_color reflcolor;
    if (maxcomponent >= dmnsn_epsilon) {
      reflcolor = dmnsn_color_mul(0.5/maxcomponent, s->color);
    } else {
      reflcolor = dmnsn_color_mul(0.25, dmnsn_white);
    }
    sphere->texture->finish.reflection = dmnsn_new_basic_reflection(
      dmnsn_black, reflcolor, 1.0
    );

    sphere->trans = dmnsn_matrix_mul(
      dmnsn_translation_matrix(s->center),
      dmnsn_scale_matrix(dmnsn_new_vector(s->radius, s->radius, s->radius))
    );

    dmnsn_array_push(scene->objects, &sphere);
  }

  dmnsn_object *plane = dmnsn_new_plane(dmnsn_y);
  plane->trans = dmnsn_translation_matrix(dmnsn_vector_mul(-4.0, dmnsn_y));
  plane->texture = dmnsn_new_texture();
  plane->texture->pigment = dmnsn_new_solid_pigment(
    dmnsn_color_from_sRGB(dmnsn_new_color(0.73, 0.90, 0.97))
  );
  plane->texture->finish.ambient = dmnsn_new_basic_ambient(
    dmnsn_color_from_sRGB(dmnsn_color_mul(0.5, dmnsn_white))
  );
  plane->texture->finish.diffuse = dmnsn_new_lambertian(
    dmnsn_sRGB_inverse_gamma(0.7)
  );
  dmnsn_array_push(scene->objects, &plane);

  dmnsn_color lcolor = dmnsn_color_mul(1.0/4.0, dmnsn_white);
  dmnsn_light *light;
  light = dmnsn_new_point_light(dmnsn_new_vector(-3.0, 0.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(-1.0, 0.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(+1.0, 0.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(+3.0, 0.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(-3.0, 5.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(-1.0, 5.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(+1.0, 5.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);
  light = dmnsn_new_point_light(dmnsn_new_vector(+3.0, 5.0, -5.0), lcolor);
  dmnsn_array_push(scene->lights, &light);

  return scene;
}

#include "dimension-internal.h"
#include <pthread.h>
#include <stdlib.h>

/** Number of children per PR-node. */
#define PRTREE_B 8
/** Number of priority leaves per pseudo-PR-node (must be 2*ndimensions). */
#define PSEUDO_B 6

/** A flat node for storing in an array for fast pre-order traversal. */
typedef struct physics_flat_prnode {
  dmnsn_bounding_box bounding_box;
  physics_sphere *object;
  size_t skip;
} physics_flat_prnode;

/** The side of the split that a node ended up on. */
typedef enum physics_prnode_location {
  PRTREE_LEAF, /**< Priority leaf. */
  PRTREE_LEFT, /**< Left child. */
  PRTREE_RIGHT /**< Right child. */
} physics_prnode_location;

/** Pseudo PR-tree node. */
typedef struct physics_prnode {
  dmnsn_bounding_box bounding_box;

  physics_sphere *object;
  struct physics_prnode *children[PRTREE_B];

  physics_prnode_location location;
} physics_prnode;

/** Construct an empty PR-node. */
static inline physics_prnode *
physics_new_prnode(void)
{
  physics_prnode *node = dmnsn_malloc(sizeof(physics_prnode));
  node->bounding_box = dmnsn_zero_bounding_box();
  node->object       = NULL;
  node->location     = PRTREE_LEFT; /* Mustn't be _LEAF */
  for (size_t i = 0; i < PRTREE_B; ++i) {
    node->children[i] = NULL;
  }
  return node;
}

/** Free a non-flat PR-tree. */
static void
physics_delete_prnode(physics_prnode *node)
{
  if (node) {
    for (size_t i = 0; i < PRTREE_B && node->children[i]; ++i) {
      physics_delete_prnode(node->children[i]);
    }
    dmnsn_free(node);
  }
}

/** Expand a node to contain the bounding box \p box. */
static void
physics_prnode_swallow(physics_prnode *node, dmnsn_bounding_box box)
{
  node->bounding_box.min = dmnsn_vector_min(node->bounding_box.min, box.min);
  node->bounding_box.max = dmnsn_vector_max(node->bounding_box.max, box.max);
}

/** Comparator types. */
enum {
  XMIN,
  YMIN,
  ZMIN,
  XMAX,
  YMAX,
  ZMAX
};

/** Get a coordinate of the bounding box of a node. */
static inline double
physics_get_coordinate(const physics_prnode * const *node, int comparator)
{
  switch (comparator) {
  case XMIN:
    return (*node)->bounding_box.min.x;
  case YMIN:
    return (*node)->bounding_box.min.y;
  case ZMIN:
    return (*node)->bounding_box.min.z;

  case XMAX:
    return -(*node)->bounding_box.max.x;
  case YMAX:
    return -(*node)->bounding_box.max.y;
  case ZMAX:
    return -(*node)->bounding_box.max.z;

  default:
    dmnsn_unreachable("Invalid comparator.");
  }
}

/* List sorting comparators */

static int
physics_xmin_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, XMIN);
  double rval = physics_get_coordinate(r, XMIN);
  return (lval > rval) - (lval < rval);
}

static int
physics_ymin_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, YMIN);
  double rval = physics_get_coordinate(r, YMIN);
  return (lval > rval) - (lval < rval);
}

static int
physics_zmin_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, ZMIN);
  double rval = physics_get_coordinate(r, ZMIN);
  return (lval > rval) - (lval < rval);
}

static int
physics_xmax_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, XMAX);
  double rval = physics_get_coordinate(r, XMAX);
  return (lval > rval) - (lval < rval);
}

static int
physics_ymax_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, YMAX);
  double rval = physics_get_coordinate(r, YMAX);
  return (lval > rval) - (lval < rval);
}

static int
physics_zmax_comp(const void *l, const void *r)
{
  double lval = physics_get_coordinate(l, ZMAX);
  double rval = physics_get_coordinate(r, ZMAX);
  return (lval > rval) - (lval < rval);
}

/** All comparators. */
static dmnsn_array_comparator_fn *const physics_comparators[PSEUDO_B] = {
  [XMIN] = physics_xmin_comp,
  [YMIN] = physics_ymin_comp,
  [ZMIN] = physics_zmin_comp,
  [XMAX] = physics_xmax_comp,
  [YMAX] = physics_ymax_comp,
  [ZMAX] = physics_zmax_comp
};

/** Add the priority leaves for this level. */
static void
physics_add_priority_leaves(dmnsn_array *sorted_leaves[PSEUDO_B],
                          dmnsn_array *new_leaves)
{
  for (size_t i = 0; i < PSEUDO_B; ++i) {
    physics_prnode *leaf = NULL;
    physics_prnode **leaves = dmnsn_array_first(sorted_leaves[i]);
    for (size_t j = 0, count = 0, size = dmnsn_array_size(sorted_leaves[i]);
         j < size && count < PRTREE_B;
         ++j)
    {
      /* Skip all the previously found extreme nodes */
      if (leaves[j]->location == PRTREE_LEAF) {
        continue;
      }

      if (!leaf) {
        leaf = physics_new_prnode();
      }
      leaves[j]->location = PRTREE_LEAF;
      leaf->children[count++] = leaves[j];
      physics_prnode_swallow(leaf, leaves[j]->bounding_box);
    }

    if (leaf) {
      dmnsn_array_push(new_leaves, &leaf);
    } else {
      return;
    }
  }
}

/** Split the sorted lists into the left and right subtrees. */
static bool
physics_split_sorted_leaves(dmnsn_array *sorted_leaves[PSEUDO_B],
                          dmnsn_array *right_sorted_leaves[PSEUDO_B],
                          size_t i)
{
  /* Get rid of the extreme nodes */
  physics_prnode **leaves = dmnsn_array_first(sorted_leaves[i]);
  size_t j, skip;
  for (j = 0, skip = 0; j < dmnsn_array_size(sorted_leaves[i]); ++j) {
    if (leaves[j]->location == PRTREE_LEAF) {
      ++skip;
    } else {
      leaves[j - skip] = leaves[j];
    }
  }
  dmnsn_array_resize(sorted_leaves[i], j - skip);

  if (dmnsn_array_size(sorted_leaves[i]) == 0) {
    return false;
  }

  /* Split the appropriate list and mark the left and right child nodes */
  right_sorted_leaves[i] = dmnsn_array_split(sorted_leaves[i]);
  DMNSN_ARRAY_FOREACH (physics_prnode **, node, sorted_leaves[i]) {
    (*node)->location = PRTREE_LEFT;
  }
  DMNSN_ARRAY_FOREACH (physics_prnode **, node, right_sorted_leaves[i]) {
    (*node)->location = PRTREE_RIGHT;
  }

  /* Split the rest of the lists */
  for (size_t j = 0; j < PSEUDO_B; ++j) {
    if (j != i) {
      right_sorted_leaves[j] = dmnsn_new_array(sizeof(physics_prnode *));

      physics_prnode **leaves = dmnsn_array_first(sorted_leaves[j]);
      size_t k, skip;
      for (k = 0, skip = 0; k < dmnsn_array_size(sorted_leaves[j]); ++k) {
        if (leaves[k]->location == PRTREE_LEAF) {
          ++skip;
        } else if (leaves[k]->location == PRTREE_RIGHT) {
          dmnsn_array_push(right_sorted_leaves[j], &leaves[k]);
          ++skip;
        } else {
          leaves[k - skip] = leaves[k];
        }
      }
      dmnsn_array_resize(sorted_leaves[j], k - skip);
    }
  }

  return true;
}

/** Recursively constructs an implicit pseudo-PR-tree and collects the priority
    leaves. */
static void
physics_priority_leaves_recursive(dmnsn_array *sorted_leaves[PSEUDO_B],
                                dmnsn_array *new_leaves,
                                int comparator)
{
  physics_add_priority_leaves(sorted_leaves, new_leaves);

  dmnsn_array *right_sorted_leaves[PSEUDO_B];
  if (physics_split_sorted_leaves(sorted_leaves, right_sorted_leaves, comparator))
  {
    physics_priority_leaves_recursive(right_sorted_leaves, new_leaves,
                                    (comparator + 1)%PSEUDO_B);
    for (size_t i = 0; i < PSEUDO_B; ++i) {
      dmnsn_delete_array(right_sorted_leaves[i]);
    }

    physics_priority_leaves_recursive(sorted_leaves, new_leaves,
                                    (comparator + 1)%PSEUDO_B);
  }
}

/** Constructs an implicit pseudo-PR-tree and returns the priority leaves. */
static dmnsn_array *
physics_priority_leaves(const dmnsn_array *leaves)
{
  dmnsn_array *sorted_leaves[PSEUDO_B];
  for (size_t i = 0; i < PSEUDO_B; ++i) {
    sorted_leaves[i] = dmnsn_array_copy(leaves);
    dmnsn_array_sort(sorted_leaves[i], physics_comparators[i]);
  }

  dmnsn_array *new_leaves = dmnsn_new_array(sizeof(physics_prnode *));
  physics_priority_leaves_recursive(sorted_leaves, new_leaves, 0);

  for (size_t i = 0; i < PSEUDO_B; ++i) {
    dmnsn_delete_array(sorted_leaves[i]);
  }

  return new_leaves;
}

static inline dmnsn_bounding_box
physics_sphere_bounding_box(const physics_sphere *sphere)
{
  dmnsn_vector rvec = dmnsn_new_vector(
    sphere->radius,
    sphere->radius,
    sphere->radius
  );
  dmnsn_bounding_box box = {
    .min = dmnsn_vector_sub(sphere->center, rvec),
    .max = dmnsn_vector_add(sphere->center, rvec),
  };
  return box;
}

/** Construct a non-flat PR-tree. */
static physics_prnode *
physics_make_prtree(const dmnsn_array *objects)
{
  if (dmnsn_array_size(objects) == 0) {
    return NULL;
  }

  /* Make the initial array of leaves */
  dmnsn_array *leaves = dmnsn_new_array(sizeof(physics_prnode *));
  DMNSN_ARRAY_FOREACH (physics_sphere *, object, objects) {
    physics_prnode *node = physics_new_prnode();
    node->bounding_box = physics_sphere_bounding_box(object);
    node->object       = object;
    dmnsn_array_push(leaves, &node);
  }

  while (dmnsn_array_size(leaves) > 1) {
    dmnsn_array *new_leaves = physics_priority_leaves(leaves);
    dmnsn_delete_array(leaves);
    leaves = new_leaves;
  }

  physics_prnode *root = *(physics_prnode **)dmnsn_array_first(leaves);
  dmnsn_delete_array(leaves);
  return root;
}

/** Recursively flatten a PR-tree into an array of flat nodes. */
static void
physics_flatten_prtree_recursive(physics_prnode *node, dmnsn_array *flat)
{
  size_t currenti = dmnsn_array_size(flat);
  dmnsn_array_resize(flat, currenti + 1);
  physics_flat_prnode *flatnode = dmnsn_array_at(flat, currenti);

  flatnode->bounding_box = node->bounding_box;
  flatnode->object       = node->object;

  for (size_t i = 0; i < PRTREE_B && node->children[i]; ++i) {
    physics_flatten_prtree_recursive(node->children[i], flat);
  }

  /* Array could have been realloc()'d somewhere else above */
  flatnode = dmnsn_array_at(flat, currenti);
  flatnode->skip = dmnsn_array_size(flat) - currenti;
}

/** Flatten a PR-tree into an array of flat nodes. */
static dmnsn_array *
physics_flatten_prtree(physics_prnode *root)
{
  dmnsn_array *flat = dmnsn_new_array(sizeof(physics_flat_prnode));
  if (root) {
    physics_flatten_prtree_recursive(root, flat);
  }
  return flat;
}

/* Construct a PR-tree from a bulk of objects */
dmnsn_array *
physics_new_prtree(const dmnsn_array *objects)
{
  physics_prnode *root = physics_make_prtree(objects);
  dmnsn_array *ret = physics_flatten_prtree(root);
  physics_delete_prnode(root);
  return ret;
}

static inline bool
physics_bounding_box_intersects(dmnsn_bounding_box box,
                                const physics_sphere *sphere)
{
  dmnsn_vector rvec = dmnsn_new_vector(
    sphere->radius,
    sphere->radius,
    sphere->radius
  );
  box.min = dmnsn_vector_max(box.min, dmnsn_vector_sub(sphere->center, rvec));
  box.max = dmnsn_vector_min(box.max, dmnsn_vector_add(sphere->center, rvec));
  return box.min.x < box.max.x
      && box.min.y < box.max.y
      && box.min.z < box.max.z;
}

void
integrate_spheres(dmnsn_array *spheres, double h)
{
  static const double g = 5.0;

  /* Inter-object collision detection */
  dmnsn_array *prtree = physics_new_prtree(spheres);
  DMNSN_ARRAY_FOREACH (physics_sphere *, s1, spheres) {
    physics_flat_prnode *node = dmnsn_array_first(prtree);
    physics_flat_prnode *last = dmnsn_array_last(prtree);
    while (node <= last) {
      physics_sphere *s2 = node->object;
      if (physics_bounding_box_intersects(node->bounding_box, s1)) {
        if (s2 && s1 != s2) {
          dmnsn_vector deltar = dmnsn_vector_sub(s1->center, s2->center);
          dmnsn_vector deltav = dmnsn_vector_sub(s1->velocity, s2->velocity);
          if (dmnsn_vector_norm(deltar) <= s1->radius + s2->radius
              && dmnsn_vector_dot(deltar, deltav) < 0.0)
          {
            dmnsn_vector x = dmnsn_vector_normalized(deltar);
            dmnsn_vector v1 = s1->velocity;
            double x1 = dmnsn_vector_dot(x, v1);
            dmnsn_vector v1x = dmnsn_vector_mul(x1, x);
            dmnsn_vector v1y = dmnsn_vector_sub(v1, v1x);

            x = dmnsn_vector_negate(x);
            dmnsn_vector v2 = s2->velocity;
            double x2 = dmnsn_vector_dot(x, v2);
            dmnsn_vector v2x = dmnsn_vector_mul(x2, x);
            dmnsn_vector v2y = dmnsn_vector_sub(v2, v2x);

            s1->velocity = dmnsn_vector_add(v2x, v1y);
            s2->velocity = dmnsn_vector_add(v1x, v2y);
          }
        }

        ++node;
      } else {
        node += node->skip;
      }
    }
  }
  dmnsn_delete_array(prtree);

  /* Floor collision detection */
  DMNSN_ARRAY_FOREACH (physics_sphere *, s, spheres) {
    if (s->center.y - s->radius <= -4.0) {
      s->velocity.y = fabs(s->velocity.y);
    }
  }

  /* Advance by the timestep */
  DMNSN_ARRAY_FOREACH (physics_sphere *, s, spheres) {
    s->center = dmnsn_vector_add(s->center, dmnsn_vector_mul(h, s->velocity));
    s->center.y -= g*h*h/2.0;
    s->velocity.y -= g*h;
  }
}

int
main()
{
  dmnsn_die_on_warnings(true);

  const double h = 1.0/25.0;
  const int nframes = 401;
  const size_t width = 1920, height = 1080;

  dmnsn_array *spheres = make_spheres();

  dmnsn_canvas *canvas = dmnsn_new_canvas(width, height);
  dmnsn_png_optimize_canvas(canvas);

  dmnsn_camera *camera = dmnsn_new_perspective_camera();
  camera->trans = dmnsn_new_matrix(
    1.7151356822004125, -0.12253122769681897, -0.2328451577118997, 3.0,
    0.0, 0.8849477555881373, -0.4656903154237999, 6.0,
    0.46776427696374845, 0.4492811682216699, 0.8537655782769662, -11.0
  );

  for (int i = 0; i < nframes; ++i) {
    if (i > 0) {
      printf("Frame %d:\t Integrating\n", i);
      static const int precision = 100;
      for (int j = 0; j < precision; ++j)
        integrate_spheres(spheres, h/precision);
    }

    printf("Frame %d:\t Rendering\n", i);
    dmnsn_scene *scene = make_scene(spheres, canvas, camera);
    dmnsn_ray_trace(scene);

    printf("Frame %d:\t Exporting\n", i);
    char fname[] = "physics00000.png";
    sprintf(fname, "physics%05d.png", i);
    FILE *image = fopen(fname, "wb");
    dmnsn_png_write_canvas(canvas, image);
    fclose(image);

    dmnsn_delete_scene(scene);
  }

  dmnsn_delete_camera(camera);
  dmnsn_delete_canvas(canvas);
  dmnsn_delete_array(spheres);
  return 0;
}