#include <array>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <random>
#include <string>
#include <vector>

#include "GL/glfw.h"
#include "GL/glu.h"

#include "chipmunk/chipmunk.h"

// Parameters
const int RESOLUTION = 10;
const double RADIUS = 0.04, WIDTH = 0.5,
  TOP = 1.5, SMALLR = 0.01, SPEED = 0.5/60.0,
  MASS = 1.0, GRAVITY = 9.8, ELASTICITY = 0.8, FRICTION = 1.0;

// Global variables
int width, height;
std::vector<std::vector<std::array<unsigned char, 3>>> image;
GLUquadric *quad;
cpSpace *space;
std::vector<cpBody *> bodies;
std::vector<std::array<float, 3>> colors;
unsigned int seed;
std::mt19937 gen;
bool paused;

void removeShape(cpBody *, cpShape *shape, void *) {
  cpSpaceRemoveShape(space, shape);
  cpShapeFree(shape);
}

void removeBody(cpSpace *, void *body, void *) {
  cpBodyEachShape((cpBody *)body, removeShape, nullptr);
  cpSpaceRemoveBody(space, (cpBody *)body);
  cpBodyFree((cpBody *)body);
}

void removeBodyPostStep(cpBody *body, void *) {
  cpSpaceAddPostStepCallback(space, removeBody, body, nullptr);
}

int quit() {
  cpSpaceEachBody(space, removeBodyPostStep, nullptr);
  cpSpaceFree(space);
  gluDeleteQuadric(quad);
  glfwTerminate();
  exit(0);
  return GL_TRUE;
}

void reset(cpSpace *, void *, void *) {
  for (auto *b : bodies)
    removeBody(space, b, nullptr);
  bodies.clear();
  gen.seed(seed);
}

void setColors() {
  int r = (double)std::min(width, height) * 0.5 * RADIUS;
  for (size_t i = 0; i < bodies.size(); ++i) {
    cpVect p = cpBodyGetPosition(bodies[i]);
    int x = std::min(std::max(std::round((p.x + 1.0) / 2.0 * width), 0.0), (double)width - 1.0);
    int y = height - 1 -
      std::min(std::max(std::round((p.y + 1.0) / 2.0 * height), 0.0), (double)height - 1.0);
    std::array<float, 3> &color = colors[i];
    color.fill(0.0f);
    float sum = 0.0f;
    for (int j = y - r; j <= y + r; ++j) {
      if (j >= 0 && j < height) {
        for (int i = x - r; i <= x + r; ++i)
          if (i >= 0 && i < width) {
            float weight =
              std::pow(r + 1 - std::abs(x - i), 2) + std::pow(r + 1 - std::abs(y - j), 2);
            sum += weight;
            color[0] += (float)image[j][i][0] / 255.0f * weight;
            color[1] += (float)image[j][i][1] / 255.0f * weight;
            color[2] += (float)image[j][i][2] / 255.0f * weight;
          }
      }
    }
    color[0] /= sum;
    color[1] /= sum;
    color[2] /= sum;
  }
}

void keypress(int key, int state) {
  if (state == GLFW_RELEASE)
    return;
  if (key == ' ') {
    setColors();
    cpSpaceAddPostStepCallback(space, reset, nullptr, nullptr);
  } else if (key == 'q')
    quit();
  else if (key == 'p')
    paused = !paused;
}

void openWindow() {
  glfwInit();
  glfwOpenWindow(width, height, 8, 8, 8, 8, 0, 0, GLFW_WINDOW);
  glfwSetWindowTitle("Falling Balls");
  glfwSwapInterval(1);
  glfwSetWindowCloseCallback(quit);
  glfwSetCharCallback(keypress);
}

void initGL() {
  glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  glEnable(GL_LIGHTING);
  glEnable(GL_LIGHT0);
  glLightModeli(GL_LIGHT_MODEL_TWO_SIDE, 1);

  GLfloat ambientLight[] = { 0.2, 0.2, 0.2, 1.0 };
  GLfloat diffuseLight[] = { 0.7, 0.7, 0.7, 1.0 };
  GLfloat specularLight[] = { 0.7, 0.7, 0.7, 1.0 };
  GLfloat position[] = { -4.5, 4.0, 4.0, 1.0 };

  glLightfv(GL_LIGHT0, GL_AMBIENT, ambientLight);
  glLightfv(GL_LIGHT0, GL_DIFFUSE, diffuseLight);
  glLightfv(GL_LIGHT0, GL_SPECULAR, specularLight);
  glLightfv(GL_LIGHT0, GL_POSITION, position);

  glDepthFunc(GL_LEQUAL);
  glEnable(GL_DEPTH_TEST);
  glEnable(GL_COLOR_MATERIAL);
  glShadeModel(GL_SMOOTH);

  quad = gluNewQuadric();
}

void addBody() {
  std::uniform_real_distribution<> r(-WIDTH, WIDTH);
  double moment = cpMomentForCircle(MASS, RADIUS, 0.0, cpvzero);
  double x, y;
  x = r(gen);
  y = TOP;
  if (!bodies.empty()) {
    for (size_t i = bodies.size(); i > 0; --i) {
      cpVect v = cpBodyGetPosition(bodies[i-1]);
      if (v.y < y - RADIUS)
        break;
      if (v.x > x - RADIUS * 2.0 && v.x < x + RADIUS * 2.0) {
        return;
      }
    }
  }
  cpBody *body = cpSpaceAddBody(space, cpBodyNew(MASS, moment));
  cpBodySetPosition(body, cpv(x, y));
  cpShape *shape = cpSpaceAddShape(space, cpCircleShapeNew(body, RADIUS, cpvzero)); // offset
  cpShapeSetElasticity(shape, ELASTICITY);
  cpShapeSetFriction(shape, FRICTION);
  bodies.push_back(body);
  if (colors.size() < bodies.size())
    colors.push_back({1.0f, 1.0f, 1.0f});
}

void initObjects() {
  space = cpSpaceNew();
  cpSpaceSetGravity(space, cpv(0, -GRAVITY));
  cpBody *static_body = cpSpaceGetStaticBody(space);
  cpShape *shape;
  shape = cpSpaceAddShape(space, cpSegmentShapeNew(static_body, cpv(-1, 1), cpv(-1, -1), SMALLR));
  cpShapeSetElasticity(shape, ELASTICITY);
  cpShapeSetFriction(shape, FRICTION);
  shape = cpSpaceAddShape(space, cpSegmentShapeNew(static_body, cpv(-1, -1), cpv(1, -1), SMALLR));
  cpShapeSetElasticity(shape, ELASTICITY);
  cpShapeSetFriction(shape, FRICTION);
  shape = cpSpaceAddShape(space, cpSegmentShapeNew(static_body, cpv(1, -1), cpv(1, 1), SMALLR));
  cpShapeSetElasticity(shape, ELASTICITY);
  cpShapeSetFriction(shape, FRICTION);
}

void updateObjects() {
  cpSpaceStep(space, SPEED);
  addBody();
}

void draw() {
  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();

  for (size_t i = 0; i < bodies.size(); ++i) {
    cpVect p = cpBodyGetPosition(bodies[i]);
    glPushMatrix();
    glTranslatef(p.x, p.y, 0.0f);
    glColor3fv(&colors[i][0]);
    gluSphere(quad, RADIUS, RESOLUTION, RESOLUTION);
    glPopMatrix();
  }

  glfwSwapBuffers();
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
}

bool readImage(const std::string &filename) {
  char code[3];
  std::ifstream f(filename);
  f.getline(code, 3);
  if (std::string(code) != "P3")
    return false;
  int r, g, b, max;
  f >> width >> height;
  f >> max;
  if (max != 255)
    return false;
  image.resize(height);
  for (int j = 0; j < height; ++j) {
    image[j].resize(width);
    for (int i = 0; i < width; ++i) {
      f >> r >> g >> b;
      image[j][i][0] = r;
      image[j][i][1] = g;
      image[j][i][2] = b;
    }
  }
  return f.good();
}

int main(int argc, char **argv) {
  if (argc != 2) {
    std::cerr << "Usage: " << argv[0] << " image.ppm" << std::endl;
    return 1;
  }

  if (!readImage(argv[1])) {
    std::cerr << "Cannot load image: " << argv[1] << std::endl;
    return 2;
  }

  std::random_device rd;
  seed = rd();
  gen.seed(seed);

  openWindow();
  initGL();
  initObjects();

  paused = false;
  while (true) {
    if (paused)
      glfwPollEvents();
    else {
      updateObjects();
      draw();
    }
  }

  return 0;
}