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

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

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)
  );

  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(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) {
        sphere s = make_sphere(x, y, z, size);
        dmnsn_array_push(spheres, &s);
      }
    }
  }

  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 (sphere *, s, spheres) {
    dmnsn_object *sph = dmnsn_new_sphere();

    sph->texture = dmnsn_new_texture();
    sph->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);
    }
    sph->texture->finish.reflection = dmnsn_new_basic_reflection(
      dmnsn_black, reflcolor, 1.0
    );

    sph->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, &sph);
  }

  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);

  scene->nthreads = 12; /* XXX */
  return scene;
}

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

  /* Inter-object collision detection */
  DMNSN_ARRAY_FOREACH (sphere *, s1, spheres) {
    DMNSN_ARRAY_FOREACH (sphere *, s2, spheres) {
      if (s1 == s2) continue;

      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);
      }
    }
  }

  /* Floor collision detection */
  DMNSN_ARRAY_FOREACH (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 (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;
}