#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <math.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SCREEN_ADDRESS 0x3C

Adafruit_SSD1306 display(
  SCREEN_WIDTH,
  SCREEN_HEIGHT,
  &Wire,
  -1,
  1700000UL,
  1700000UL
);

// Put the generated protein_model.h in this sketch folder.
#include "protein_model.h"

int touch_pins[6] = {3, 4, 2, 27, 1, 26};

// const int TOUCH0_PIN = 3;
// const int TOUCH1_PIN = 4;
const int TOUCH_THRESHOLD = 6;
const int TOUCH_MAX = 200;
const float X_ROTATION_SPEED = 0.08;
const float Y_ROTATION_SPEED = 0.08;

// scale up to take up the whole screen
const float MODEL_SCALE = 1.25;

float angleX = 0.0;
float angleY = 0.0;

struct ScreenSegment {
  int x1;
  int y1;
  int x2;
  int y2;
  float depth;
};

int measureTouch(int pin) {
  int count = 0;

  pinMode(pin, OUTPUT);
  digitalWrite(pin, LOW);
  delayMicroseconds(25);

  pinMode(pin, INPUT_PULLUP);
  while (!digitalRead(pin) && count < TOUCH_MAX) {
    count++;
  }

  return count;
}

Point3D rotateAroundX(Point3D point, float angle) {
  Point3D result;
  result.x = point.x;
  result.y = point.y * cos(angle) - point.z * sin(angle);
  result.z = point.y * sin(angle) + point.z * cos(angle);
  return result;
}

Point3D rotateAroundY(Point3D point, float angle) {
  Point3D result;
  result.x = point.x * cos(angle) + point.z * sin(angle);
  result.y = point.y;
  result.z = -point.x * sin(angle) + point.z * cos(angle);
  return result;
}

Point3D rotatePoint(Point3D point) {
  return rotateAroundY(rotateAroundX(point, angleX), angleY);
}

// Draw two pixels on, then two pixels off along a line.
void drawDottedLine(int x1, int y1, int x2, int y2) {
  int dx = abs(x2 - x1);
  int sx = x1 < x2 ? 1 : -1;
  int dy = -abs(y2 - y1);
  int sy = y1 < y2 ? 1 : -1;
  int error = dx + dy;
  int step = 0;

  while (true) {
    if ((step % 4) < 2) {
      display.drawPixel(x1, y1, SSD1306_WHITE);
    }

    if (x1 == x2 && y1 == y2) break;

    int doubledError = 2 * error;
    if (doubledError >= dy) {
      error += dy;
      x1 += sx;
    }
    if (doubledError <= dx) {      
      error += dx;
      y1 += sy;
    }
    step++;
  }
}

void drawThickLine(int x1, int y1, int x2, int y2) {
  display.drawLine(x1, y1, x2, y2, SSD1306_WHITE);

  // Offset the second line across the segment rather than along it.
  if (abs(x2 - x1) >= abs(y2 - y1)) {
    display.drawLine(x1, y1 + 1, x2, y2 + 1, SSD1306_WHITE);
  } else {
    display.drawLine(x1 + 1, y1, x2 + 1, y2, SSD1306_WHITE);
  }
}

void sortBackToFront(ScreenSegment segments[], int count) {
  // Insertion sort is adequate for the small number of protein segments.
  for (int i = 1; i < count; i++) {
    ScreenSegment current = segments[i];
    int j = i - 1;

    while (j >= 0 && segments[j].depth > current.depth) {
      segments[j + 1] = segments[j];
      j--;
    }
    segments[j + 1] = current;
  }
}

void setup() {
  Serial.begin(115200);

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    while (true) {
      // Stop if the OLED could not be initialized.
    }
  }
}

void loop() {
  bool touchUp = measureTouch(touch_pins[5]) > TOUCH_THRESHOLD;
  bool touchDown = measureTouch(touch_pins[2]) > TOUCH_THRESHOLD;

  if (touchDown) angleX += X_ROTATION_SPEED;
  if (touchUp) angleX -= X_ROTATION_SPEED;

  bool touchLeft = measureTouch(touch_pins[3]) > TOUCH_THRESHOLD;
  bool touchRight = measureTouch(touch_pins[4]) > TOUCH_THRESHOLD;

  if (touchRight) angleY += Y_ROTATION_SPEED;
  if (touchLeft) angleY -= Y_ROTATION_SPEED;

  // // Y rotation continues slowly without user input.
  // angleY += Y_ROTATION_SPEED;

  if (angleX > 2 * PI) angleX -= 2 * PI;
  if (angleX < -2 * PI) angleX += 2 * PI;
  if (angleY > 2 * PI) angleY -= 2 * PI;
  if (angleY < -2 * PI) angleY += 2 * PI;

  display.clearDisplay();

  ScreenSegment segments[NUM_POINTS - 1];
  float minimumDepth = 1000000.0;
  float maximumDepth = -1000000.0;

  for (int i = 0; i < NUM_POINTS - 1; i++) {
    Point3D a = rotatePoint(protein[i]);
    Point3D b = rotatePoint(protein[i + 1]);

    segments[i].x1 = 64 + (int)(a.x * MODEL_SCALE);
    segments[i].y1 = 32 - (int)(a.y * MODEL_SCALE);
    segments[i].x2 = 64 + (int)(b.x * MODEL_SCALE);
    segments[i].y2 = 32 - (int)(b.y * MODEL_SCALE);
    segments[i].depth = (a.z + b.z) * 0.5;

    if (segments[i].depth < minimumDepth) minimumDepth = segments[i].depth;
    if (segments[i].depth > maximumDepth) maximumDepth = segments[i].depth;
  }

  sortBackToFront(segments, NUM_POINTS - 1);

  float depthRange = maximumDepth - minimumDepth;
  if (depthRange < 0.001) depthRange = 1.0;

  // Distant segments are drawn first, so nearby ones visually overlap them.
  for (int i = 0; i < NUM_POINTS - 1; i++) {
    float relativeDepth =
      (segments[i].depth - minimumDepth) / depthRange;

    if (relativeDepth < 0.33) {
      drawDottedLine(
        segments[i].x1, segments[i].y1,
        segments[i].x2, segments[i].y2
      );
    } else if (relativeDepth < 0.67) {
      display.drawLine(
        segments[i].x1, segments[i].y1,
        segments[i].x2, segments[i].y2,
        SSD1306_WHITE
      );
    } else {
      drawThickLine(
        segments[i].x1, segments[i].y1,
        segments[i].x2, segments[i].y2
      );
    }
  }

  display.display();
  delay(30);
}