#include <joyos.h>
#include <kern/lock.h>
#include <lib/pid.h>
#include <math.h>

#define FORWARD_VEL 104
#define LSB_US_PER_DEG 1468056
#define MULTIPLIER -1.1
#define HEADING_RANGE 20
#define DISTANCE_RANGE 500
#define ENCODER_CONST 50 // Lower values cause overshoot. Higher values cause undershoot.

struct lock data_lock;
struct pid_controller headingPID;
float turn = 0;
float dist = 0;
int dir = 1;
int pidmode = 0; // 1 = turn. 2 = adjust.
int target = 0;
int tarType = 0; // 1 = golf, 2 tennis, 3 = rift bomb.

int bomb = 0;

int spd = 0;

/**
 * Limits the velocity to an integer between 0 and 255.
 */
int16_t limitVel(float vel) {
  if (vel < 0) {
    return 0;
  }
  if (vel > 255) {
    return 255;
  }
  return (int16_t)vel;
}

/**
 * Helper function that determines the distance to a point.
 */
float distTo(int x, int y) {
  int32_t distx = x - objects[0].x;
  int32_t disty = y - objects[0].y;
  return sqrt(distx * distx + disty * disty);
}

/**
 * Helper function that normalizes angles to -180 < theta <= 180.
 */
float normalize(float angle) {
  angle = fmod(angle, 360);
  if (angle > 180) {
    angle = angle - 360;
  }
  return angle;
}

/**
 * Gets the heading. Returns heading relative to intended direction = 0.
 */
float gyro_get_degrees_with_offset() {
  // Goal: 0.
  float angle = turn - gyro_get_degrees();
  angle = normalize(angle);
  if (dir == -1) {
    if (angle < -90) {
      angle = -180 - angle;
    }
    else {
      angle = 180 - angle;
    }
  }
  // printf("Angle: %f\nDirection: %d\nGyro: %f\nTurn: %f\n\n", angle, dir, gyro_get_degrees(), turn);
  return angle;
}

/**
 * Sets the velocities, based on data from the nav system.
 */
void setVelocities(float correctionAmount) {
  int localMode = pidmode;
  int localDir = dir;
  // printf("Movement:\n          Dir: %d\n          Mod: %d\n          Cor: %f\n\n", localDir, localMode, correctionAmount);
  if (localMode == 0) {
    // Don't move, return.
    return;
  }
  if (localMode == 1) {
    /* @TODO
    if (correctionAmount < 0) {
      motor_set_vel(0, localDir * limitVel((FORWARD_VEL * 0.75 - correctionAmount) * MULTIPLIER));
      motor_set_vel(2, -1 * localDir * limitVel(FORWARD_VEL * 0.75 - correctionAmount));
    }
    else {
      motor_set_vel(0, -1 * localDir * limitVel((FORWARD_VEL * 0.75 + correctionAmount) * MULTIPLIER));
      motor_set_vel(2, localDir * limitVel(FORWARD_VEL * 0.75 + correctionAmount));
    }*/
  }
  if (localMode == 2) {
    int fv = spd;
    if (fv == 0) {
      fv = FORWARD_VEL;
    }
    motor_set_vel(0, localDir * limitVel((fv - correctionAmount)) * MULTIPLIER);
    motor_set_vel(2, localDir * limitVel(fv + correctionAmount));
    // printf("Setting velocities, right: %d, left: %d", localDir * limitVel(FORWARD_VEL - correctionAmount), localDir * limitVel(FORWARD_VEL + correctionAmount));
  }
}

/**
 * Determine which ball to go for. Go for bigger balls first.
 */
int goFor() {
  // Order by radius, then by distance. Radius difference of <= 2 doesn't matter.
  int num = 0;
  int curType = 0; // 0 = none. 1 = Golf. 2 = Tennis. 3 = Bomb.
  int tmpType = 0;
  float curDist = 0;
  float tmpDist = 0;
  int i;
  copy_objects();
  for (i = 2; i <= 31; i++) {
    if (objects[i].id == 255) {
      // Only go for balls on the same side as the robot.
      if (((objects[0].y > 0) && (objects[i].y > 0)) || ((objects[0].y < 0) && (objects[i].y < 0))) {
        if (objects[i].radius < 8) {
          tmpType = 1;
        }
        else {
          if (objects[i].hue <= 1 || objects[i].hue >= 13) {
            // Pink.
            tmpType = 2;
          }
          else if (objects[i].hue >= 2 && objects[i].hue <= 5) {
            // Green-yellow.
            tmpType = 2;
          }
          else {
            // Other.
            tmpType = 3;
            bomb = i;
          }
          if (bomb == i) {
            tmpType = 3;
          }
        }
        tmpDist = distTo(objects[i].x, objects[i].y);
        if (tmpType > curType) {
          num = i;
          curDist = tmpDist;
          curType = tmpType;
        }
        else if (tmpType == curType) {
          if (tmpDist < curDist) {
            num = i;
            curDist = tmpDist;
            curType = tmpType;
          }
        }
      }
    }
  }
  tarType = curType;
  return num;
}

void turnToHeading(float heading) {
  /*
  @TODO
  float kP = 0.4; //proportional gain
  float kI = 0.0; //integral gain
  float kD = 0.0; //derivative gain

  // printf("Acquiring lock\n");
  acquire(&data_lock);
  pidmode = 1;
  // printf("Initing pid\n");
  init_pid( &headingPID,
            kP,
            kI,
            kD,
            &gyro_get_degrees_with_offset,
            &setVelocities);
  turn = heading;
  headingPID.enabled = true;
  headingPID.goal = 0;
  release(&data_lock);
  */
  acquire(&data_lock);
  turn = heading;
  release(&data_lock);
  float curHead = gyro_get_degrees_with_offset();
  int index = 0;
  while ((curHead < -HEADING_RANGE || curHead > HEADING_RANGE) && (index < 100)) {
    // Keep turning until heading in roughly the correct direction.
    if (curHead > 0) {
      motor_set_vel(0, FORWARD_VEL * 0.75 * MULTIPLIER);
      motor_set_vel(2, -FORWARD_VEL * 0.75);
    }
    else {
      motor_set_vel(0, -FORWARD_VEL * 0.75 * MULTIPLIER);
      motor_set_vel(2, FORWARD_VEL * 0.75);
    }
    curHead = gyro_get_degrees_with_offset();
    index++;
    pause(20);
  }
  acquire(&data_lock);
  pidmode = 0;
  release(&data_lock);
  motor_brake(0);
  motor_brake(2);
}

void turnToPoint(int32_t pointX, int32_t pointY) {
  int posx = objects[0].x;
  int posy = objects[0].y;
  int32_t distx = pointX - posx;
  int32_t disty = pointY - posy;
  float heading = atan2(disty, distx) * 180 / M_PI;
  turnToHeading(heading);
}

void moveToPoint(int32_t pointX, int32_t pointY) {
  copy_objects();
  float local_dist = distTo(pointX, pointY);
  int posx = objects[0].x;
  int posy = objects[0].y;
  int32_t distx = pointX - posx;
  int32_t disty = pointY - posy;
  float heading = atan2(disty, distx) * 180 / M_PI;

  float kP = 1.5; //proportional gain
  float kI = 0.0; //integral gain
  float kD = 0.0; //derivative gain

  encoder_reset(24);
  encoder_reset(25);

  acquire(&data_lock);
  pidmode = 2;
  init_pid( &headingPID,
            kP,
            kI,
            kD,
            &gyro_get_degrees_with_offset,
            &setVelocities);
  headingPID.enabled = true;
  headingPID.goal = 0;
  dist = local_dist;
  turn = heading;
  release(&data_lock);
  int index = 0;
  while (((dir == -1) || (local_dist - DISTANCE_RANGE > (encoder_read(24) + encoder_read(25)) * ENCODER_CONST)) && (index < 150)) {
    // Keep going until almost there.
    pause(20);
    index++;
  }
  acquire(&data_lock);
  pidmode = 0;
  release(&data_lock);
}

void squeeze() {
  servo_set_pos(2, 511);
  pause(500);
}

void unsqueeze() {
  servo_set_pos(2, 127);
  pause(500);
}

void scoop() {
  servo_set_pos(0, 127);
  servo_set_pos(1, 383);
  pause(1500);
}

void unscoop() {
  servo_set_pos(0, 447);
  servo_set_pos(1, 64);
  pause(500);
}

/**
 * Recalibrates the gyro heading to the data from the system.
 */
void recalibrate() {
  acquire(&data_lock);
  copy_objects();
  motor_brake(0);
  motor_brake(2);
  pause(500);
  float t = objects[0].theta * 45.0 / 512.0;
  gyro_set_degrees(t);
  release(&data_lock);
}

void releaseBomb() {
  if (objects[0].y > 0) {
    turnToPoint(-1907, 1365);
    pause(200);
    moveToPoint(-1907, 1365);
    motor_brake(0);
    motor_brake(2);
    pause(200);
    turnToHeading(-90);
    pause(200);
    copy_objects();
    moveToPoint(objects[0].x, 341);
    motor_brake(0);
    motor_brake(2);
    pause(200);
    turnToHeading(0);
    pause(200);
    acquire(&data_lock);
    dir = -1;
    spd = 72;
    release(&data_lock);
    moveToPoint(-2047, 720);
  }
  else {
    turnToPoint(1707, -1365);
    pause(200);
    moveToPoint(1707, -1365);
    motor_brake(0);
    motor_brake(2);
    pause(200);
    turnToHeading(90);
    pause(200);
    copy_objects();
    moveToPoint(objects[0].x, -341);
    motor_brake(0);
    motor_brake(2);
    pause(200);
    turnToHeading(180);
    pause(200);
    acquire(&data_lock);
    dir = -1;
    spd = 72;
    release(&data_lock);
    moveToPoint(2047, -720);
  }
  /*motor_set_vel(0, -60 * MULTIPLIER);
  motor_set_vel(2, -60);
  motor_set_vel(3, 150);
  pause(1500);*/
  motor_set_vel(3, 128);
  motor_brake(0);
  motor_brake(2);
  acquire(&data_lock);
  dir = 1;
  spd = 0;
  release(&data_lock);
  motor_set_vel(0, 30 * MULTIPLIER);
  motor_set_vel(2, 30);
  pause(15000);
  motor_brake(3);

}

void golfBallSweep(int x, int y) {
  copy_objects();
  turnToPoint(x, y);
  moveToPoint(x, y);
  pause(333);
  motor_brake(0);
  motor_brake(2);
  motor_set_vel(0, -FORWARD_VEL * MULTIPLIER);
  motor_set_vel(2, -FORWARD_VEL);
  pause(1750);
  motor_brake(0);
  motor_brake(2);
}

/**
 * Main navigation loop. Hanadles high-level navigation, such as where to go,
 * where to turn, and updating the pid controller.
 */
int navigate() {
  int localMode;
  int i;
  while (true) {
    acquire(&data_lock);
    localMode = pidmode;
    release(&data_lock);
    if (localMode != 0) {
      copy_objects();
      acquire(&data_lock);
      update_pid(&headingPID);
      release(&data_lock);
    }
    yield();
    pause(20);
    i++;
    if (i % 50 == 0) {
      // Check for rift bombs.
      goFor();
    }
  }
  return 0;
}

/**
 * Main thread. Handles high level code, directing the navigation thread and
 * highest-level strategy.
 */
int umain() {
  init_lock(&data_lock, "Data_Lock");
  // Goes for the 4 balls in front of it
  /*motor_set_vel(0, FORWARD_VEL * MULTIPLIER * 0.75);
  motor_set_vel(02, FORWARD_VEL * 1.1 * 0.75);
  pause(1200);
  motor_set_vel(0, -FORWARD_VEL * MULTIPLIER * 0.75);
  motor_set_vel(2, -FORWARD_VEL * 1.1 * 0.75);
  pause(600);
  motor_brake(2);
  motor_brake(0);
  */
  create_thread(&navigate, STACK_DEFAULT, 0, "nav_thread");
  copy_objects();
  recalibrate();
  if (objects[0].y > 0) {
    turnToHeading(-90);
  }
  else {
    turnToHeading(90);
  }
  unscoop();

  recalibrate();
  copy_objects();
  if (objects[0].y > 0) {
    golfBallSweep(1700, 256);
    pause(200);
    golfBallSweep(1200, 256);
  }
  else {
    golfBallSweep(-1700, -256);
    pause(200);
    golfBallSweep(-1200, -256);
  }
  // Decided not to release the bomb. Oh well.
  // releaseBomb();
  while (true) {
    recalibrate();
    int tar = goFor();
    bomb = 0;
    copy_objects();
    acquire(&data_lock);
    target = tar;
    release(&data_lock);
    turnToPoint(objects[tar].x, objects[tar].y);
    copy_objects();
    moveToPoint(objects[tar].x, objects[tar].y);
    pause(200);
    squeeze();
    motor_brake(0);
    motor_brake(2);
    if (tarType == 1) {
      // Golf ball.
      if (objects[0].y < 0) {
        turnToHeading(90);
      }
      else {
        turnToHeading(-90);
      }
      unsqueeze();
      copy_objects();
      moveToPoint(objects[0].x, 0);
      motor_set_vel(0, -FORWARD_VEL * MULTIPLIER);
      motor_set_vel(2, -FORWARD_VEL);
    }
    else {
      // Everything else.
      motor_set_vel(0, -FORWARD_VEL * MULTIPLIER);
      motor_set_vel(2, -FORWARD_VEL);
      pause(600);
      motor_brake(0);
      motor_brake(2);
      copy_objects();
      if (objects[0].y < 0) {
        turnToHeading(-90);
      }
      else {
        turnToHeading(90);
      }
      /*copy_objects();
      int i = 0;
      while (abs(objects[0].y) > BACKUP_TO && i < 2) {*/
        acquire(&data_lock);
        dir = -1;
        release(&data_lock);
        copy_objects();
        moveToPoint(objects[0].x, 0);
        //i++;
  //  }
      motor_brake(0);
      motor_brake(2);
      acquire(&data_lock);
      dir = 1;
      release(&data_lock);
      scoop();
      unscoop();
      unsqueeze();
      motor_set_vel(0, FORWARD_VEL * MULTIPLIER);
      motor_set_vel(2, FORWARD_VEL);
    }
    pause(750);
    motor_brake(0);
    motor_brake(2);
  }
  return 0;
}

/**
 * Setup function. Initializes and recalibrates the gyro.
 */
int usetup() {
  extern volatile uint8_t robot_id;
  robot_id = 9;

  // Initialize the gyroscope.
  gyro_init(11, LSB_US_PER_DEG, 5000L);
  recalibrate();

  // Debug code:
  /*copy_objects();
  int i;
  for (i = 2; i <= 31; i++) {
    if (objects[i].id == 255) {
      // Only go for balls on the same side as the robot.
      if (((objects[0].y > 0) && (objects[i].y > 0)) || ((objects[0].y < 0) && (objects[i].y < 0))) {
        printf("%d: ", i);
        //printf("%d:    X: %d    Y: %d    Dist: %f    Rad: %d     Hue: %d    Sat:    %d\n", i, objects[i].x, objects[i].y, distTo(objects[i].x, objects[i].y), objects[i].radius, objects[i].hue, objects[i].saturation);
      }
    }
  }*/

  // create_thread(&navigate, STACK_DEFAULT, 0, "nav_thread");
  // turnToHeading(-90);

  // dir = -1;
  // create_thread(&navigate, STACK_DEFAULT, 0, "nav_thread");
  // copy_objects();
  // int sign = abs(objects[0].y) / objects[0].y;
  // moveToPoint(sign * 1024, sign * 1024);

  // squeeze();
  // scoop();
  // unscoop();
  // unsqueeze();

  /*motor_set_vel(0, FORWARD_VEL * MULTIPLIER);
  motor_set_vel(2, FORWARD_VEL);
  pause(2000);
  motor_brake(0);
  motor_brake(2);*/

  return 0;
}