#include <Servo.h>
#include <SPI.h>
#include <nRF24L01.h>
#include <RF24.h>
#include <Wire.h>

/* Pin definitions */
#define LEFT_DRIVE_MOTOR_CONTROLPIN_1 9
#define LEFT_DRIVE_MOTOR_CONTROLPIN_2 10
#define LEFT_DRIVE_MOTOR_SPEEDPIN 8

#define RIGHT_DRIVE_MOTOR_CONTROLPIN_1 12
#define RIGHT_DRIVE_MOTOR_CONTROLPIN_2 11
#define RIGHT_DRIVE_MOTOR_SPEEDPIN 13

#define WATER_SERVO_CONTROLPIN_1 3
#define WATER_SERVO_CONTROLPIN_2 4
#define WATER_SERVO_SPEEDPIN 2

#define FIRE_SERVO_CONTROLPIN_1 6
#define FIRE_SERVO_CONTROLPIN_2 5
#define FIRE_SERVO_SPEEDPIN 7

#define TOWER_X_PIN 26
#define TOWER_Y_PIN 25

#define PILOT_FLAME 24

#define PRESSURE_SENSOR_PIN A4
#define COMPRESSOR_RELAY_PIN 41

#define RADIO_CE_PIN 49
#define RADIO_CSN_PIN 53
/*
SCK  52
MOSI 51
MISO 50
*/

#define RGB_LED_RED_PIN A0
#define RGB_LED_GREEN_PIN A1
#define RGB_LED_BLUE_PIN A2


/* variables used by i2c protcol */
#define SLAVE_ADDRESS 0x04
#define ST_IDLE 0x1
#define ST_HEAD1_OK 0x2
#define ST_HEAD2_OK 0x3
#define ST_HEADER_OK 0x4
#define ST_X_READ 0x5

#define CAMERA_COORDINATES_X_THRESHOLD 40
#define CAMERA_COORDINATES_Y_THRESHOLD 10
#define AUTO_TURN_SPEED 255
#define AUTO_DRIVE_SPEED 255
#define AUTO_EXTINGUISH_THRESHOLD 40
#define AUTO_TOWER_Y_MOVE_STEPS 127
#define AUTO_EXTINGUISH_TIME 2000

//enum of modes
enum Mode {
  WATERMODE,
  FIREMODE,
  AUTOMODE
};

/*
  Class to control the motors
*/
class MotorControl {
  private:
    int controlPin1, controlPin2, speedPin;
  public:
    MotorControl(int p1, int p2, int p3) {
      controlPin1 = p1;
      controlPin2 = p2;
      speedPin = p3;

      pinMode(controlPin1, OUTPUT);
      pinMode(controlPin2, OUTPUT);
      pinMode(speedPin, OUTPUT);
    }
    void setSpeed(int s) {
      //if speed is set to 0, no further calculations is needed
      if(s == 0) {
        analogWrite(speedPin, 0);
        return;
      }
      //set the control pins to (0, 1) or (1, 0) based on the direction
      digitalWrite(controlPin1, s > 0);
      digitalWrite(controlPin2, s < 0);

      s = abs(s); //direction is handled by the control pins
      s = map(s, 0, 512, 60, 255); //map the speed value from 0-512 to 60-255. 60 is the minimum speed for the motors to start
      s = constrain(s, 60, 255); //constrain the variable
      analogWrite(speedPin, s); //write it to the speed control
    }
};

//The lego servos use a similar control to the motors
class BinaryMotorControl : protected MotorControl {
  private:
    //two states given
    int states[2];
  public:
    //initialize the servo with the states
    BinaryMotorControl(int s1, int s2, int p1, int p2, int p3) : MotorControl(p1, p2, p3) {
      states[0] = s1;
      states[1] = s2;
    }
    void setState(boolean state) {
      setSpeed(states[state]);
    }
};

/* Class to control the servos */
class JoystickServo {
  private:
    Servo s;
    //lower and upper range, that the servo will stay within
    int upperRange, lowerRange, resetPos;
    //position of the servo, float so its possible to 'move' less than a step
    float pos;
  public:
    JoystickServo(int pin, int lr, int ur, int start) {
      //initialize variables and put the servo in its starting position
      s.attach(pin);
      resetPos = start;
      pos = start;
      s.write(start);
      lowerRange = lr;
      upperRange = ur;
    }
    //move the servo, 512 being 1 step
    void move(int value) {
      pos += value/512.0;
      if(pos < lowerRange) pos = lowerRange;
      else if(pos > upperRange) pos = upperRange;
      s.write(pos);
    }
    int getPos() {
      //get the current position of the servo
      return pos;
    }
    void reset() {
      //put the servo back to its starting position
      pos = resetPos;
      s.write(pos);
    }
};

/* Servo controller with two states */
class ToggleServo {
  private:
    int states[2];
    //variables to slow down the movement of the servo
    float pos;
    boolean moving;
    Servo servo;
  public:
    ToggleServo(int pin, int s1, int s2) {
      servo.attach(pin);
      servo.write(s1);
      moving = 0;
      states[0] = s1;
      states[1] = s2;
    }
    void setState(boolean s) {
      //slow down only one direction, reverting to start position is instant
      if(s) {
        moving = true;
      } else {
        servo.write(states[0]);
        pos = states[0];
      }
    }
    void tick() {
      //if the servo is supposed to move
      if(moving) {
        //if the servo has not yet reached its goal position
        if(pos > states[1]) pos -= 0.8;
        else {
          //if it has, stop 'moving'
          pos = states[1];
          moving = 0;
        }
        //write the position to the servo
        servo.write(pos);
      }
    }
    boolean isOn() {
      //check wether the servo is on or off
      return pos == states[1];
    }
};

/*
  Class to maintain the pressure in the tanks. Fully automatic
*/
class PressureSustainer {
  private:
    float upper, lower, bar, avg;
    int count, sensorPin, relayPin;
  public:
    PressureSustainer(float b, float t, int s, int r) {
      upper = b;
      lower = b - t;
      bar = 0.0;
      setState(0);
      count = 10;
      avg = 0.0;

      sensorPin = s;
      relayPin = r;
      pinMode(r, OUTPUT);
      digitalWrite(r, 1);
    }
    void tick() {
      //map the values from the sensor (0.5V - 4.5V) to the measure range (0bar - 12bar)
      bar = map(analogRead(sensorPin), 102, 921, 0, 1200) / 100.0;

      //count the 10 last measurements
      avg += bar;
      count--;
      //when 10 measurements are taken
      if(count <= 0) {
        avg /= 10.0; //calculate the average
        if(avg < lower) setState(1); //if the average is lower than the goal - threshold, start compressor
        else if(avg > upper) setState(0); //if its higher than goal, stop the compressor
        count = 10; //start count over
        avg = 0.0;
      }
    }
    void setState(bool state) {
      //turn compressor on/off
      digitalWrite(relayPin, !state);
    }
    float getBar() {
      //return the current bar value
      return bar;
    }
    int getValue() {
      //return value of the sensor
      return analogRead(sensorPin);
    }
};

//function to decompress the data received
void byteArrToIntArr(byte arr[], int byteSize, int keep[], int data[], int dataSize) {
  uint64_t temp = 0;
  for(int i = 0; i < byteSize; i++) temp = (temp << 8) | arr[i]; //push the byte into the temporary 64b variable
  for(int i = dataSize - 1; i >= 0; i--) {
    data[i] = (int)(temp & ((1 << keep[i]) - 1)); //take x number of bits and put them into its own variable
    temp >>= keep[i];
  }
}


//pipe for communication
const uint64_t pipe = 0xF0F0F0F0AALL;


//pointers as not all arduino stuff are ready yet (pinMode eg.)
ToggleServo *pilotFlame;
JoystickServo *towerX;
JoystickServo *towerY;
MotorControl *leftDrive;
MotorControl *rightDrive;
BinaryMotorControl *waterValve;
BinaryMotorControl *fireValve;

PressureSustainer *pressureSustainer;

RF24 radio(RADIO_CE_PIN, RADIO_CSN_PIN);

//communication variables
int bytesReceived = 6;
byte arrayReceived[6];

//the data received, how many bits to store and what they are
int data[8];
int bitsToKeep[] = {10, 10, 1, 10, 10, 1, 2, 1};
enum Data {
  J1X,
  J1Y,
  PILOTFLAME,
  J2X,
  J2Y,
  TRIGGER,
  MODE
};

//initializing variables
int leftMotorSpeed, rightMotorSpeed, x, y;
int state = ST_IDLE; //I2C state
int cameraCoordinates[2];
bool newCoordinatesReceived, flameFound;
Mode mode;

//array of the led pins, their colours in relation to the modes (water, fire, auto)
byte leds[] = {RGB_LED_BLUE_PIN, RGB_LED_RED_PIN, RGB_LED_GREEN_PIN};




void setup() {
  //start up radio and set payload size
  radio.begin();
  radio.setPayloadSize(6);
  radio.openReadingPipe(1, pipe);
  radio.startListening();

  //start up I2C communication
  Wire.begin(SLAVE_ADDRESS);
  Wire.onReceive(receiveData);
  newCoordinatesReceived = 0;
  flameFound = 0;

  //initialize classes with given parameters
  pilotFlame = new ToggleServo(PILOT_FLAME, 108, 55);
  towerX = new JoystickServo(TOWER_X_PIN, 15, 120, 84);
  towerY = new JoystickServo(TOWER_Y_PIN, 10, 110, 50);
  leftDrive = new MotorControl(LEFT_DRIVE_MOTOR_CONTROLPIN_1, LEFT_DRIVE_MOTOR_CONTROLPIN_2, LEFT_DRIVE_MOTOR_SPEEDPIN);
  rightDrive = new MotorControl(RIGHT_DRIVE_MOTOR_CONTROLPIN_1, RIGHT_DRIVE_MOTOR_CONTROLPIN_2, RIGHT_DRIVE_MOTOR_SPEEDPIN);
  waterValve = new BinaryMotorControl(0, 150, WATER_SERVO_CONTROLPIN_1, WATER_SERVO_CONTROLPIN_2, WATER_SERVO_SPEEDPIN);
  fireValve = new BinaryMotorControl(0, 150, FIRE_SERVO_CONTROLPIN_1, FIRE_SERVO_CONTROLPIN_2, FIRE_SERVO_SPEEDPIN);
  pressureSustainer = new PressureSustainer(3, 0.3, PRESSURE_SENSOR_PIN, COMPRESSOR_RELAY_PIN);

  //set the mode for the led pins
  for(int i = 0; i < sizeof(leds); i++) pinMode(leds[i], OUTPUT);
}

void loop() {
  mode = AUTOMODE;
  //set automode as standard
  while(radio.available()) {
    //if radio is received, read it
    radio.read(arrayReceived, bytesReceived);
    //translate the byte string received to an int array
    byteArrToIntArr(arrayReceived, bytesReceived, bitsToKeep, data, 8);

    //subtract 511 to obtain position data with 0 as origa
    data[J1X] -= 511;
    data[J1Y] -= 511;
    data[J2X] -= 511;
    data[J2Y] -= 511;

    //set the mode from the radio signal
    mode = data[MODE];
    //if in automatic mode, break out of manual control loop
    if(mode == AUTOMODE) break;


    if(mode == WATERMODE) {
      //if in watermode, ensure firevalve is closed
      fireValve->setState(0);
      //ensure pilot flame is off
      pilotFlame->setState(0);
      //set the status of the watervalve to the value of the trigger send over radio
      waterValve->setState(data[TRIGGER]);
    } else if(mode == FIREMODE) {
      //if in firemode, ensure water valve is closed
      waterValve->setState(0);
      //turn on pilot flame if depending on radio signal
      pilotFlame->setState(data[PILOTFLAME]);
      if(pilotFlame->isOn()) {
        //if the pilot flame is on, set firevalve to the status of the trigger signal
        fireValve->setState(data[TRIGGER]);
      }
    }

    //variables for computation of drive engine speed
    x = data[J1X];
    y = data[J1Y];

    //set both engines to have speed depending on y-axis on the joystick (-512, 512)
    leftMotorSpeed = y;
    rightMotorSpeed = y;

    //adding the turn value (x-axis) to the side in question
    if(x > 0) leftMotorSpeed += x;
    else rightMotorSpeed -= x;

    //setting the drive speeds
    leftDrive->setSpeed(leftMotorSpeed);
    rightDrive->setSpeed(rightMotorSpeed);

    //moving the tower servos depending on input from ther ight joystick
    towerX->move(-data[J2X]);
    towerY->move(data[J2Y]);
  }
  /* Make sure the correct led is lit, depending on state */
  for(int i = 0; i < sizeof(leds); i++) {
    if(i == mode) digitalWrite(leds[i], 1);
    else digitalWrite(leds[i], 0);
  }

  if(mode == AUTOMODE) {
    //if pilot flame is on, turn it off
    if(pilotFlame->getState()) pilotFlame->setState(0);

    //if a flame is directly in front
    if(flameFound) {
      if(newCoordinatesReceived)   else {

        //the flame is extinguished, set that state, turn of water valve and reset the tower position
        flameFound = 0;
        waterValve->setState(0);
        towerX->reset();
        towerY->reset();
      }
    } else if(!newCoordinatesReceived) {
      //no new coordinates from the camera, rotate on the spot (until flame is found)
      leftDrive->setSpeed(-AUTO_TURN_SPEED);
      rightDrive->setSpeed(AUTO_TURN_SPEED);
    } else {
      //the robot sees the flame
      if(abs(cameraCoordinates[1]) > CAMERA_COORDINATES_Y_THRESHOLD) {
        //if the camera is out of the Y-axis thershold, move the tower along that axis
        towerY->move(AUTO_TOWER_Y_MOVE_STEPS * -cameraCoordinates[1]/100);
      }
      if(abs(cameraCoordinates[0]) > CAMERA_COORDINATES_X_THRESHOLD) {
        //if the flame is outside the X-axis thershold, rotate the robot towards it
        int modifier = cameraCoordinates[0] < 0 ? -1 : 1;
        leftDrive->setSpeed(AUTO_TURN_SPEED * modifier);
        rightDrive->setSpeed(-AUTO_TURN_SPEED * modifier);

      } else {
        //the flame is within the x-axis threshold
        if(towerY->getPos() > AUTO_EXTINGUISH_THRESHOLD) {
          //check the position of the tower along y axis, if its above the the set 'angle' (servo position), drive forward
          leftDrive->setSpeed(AUTO_DRIVE_SPEED);
          rightDrive->setSpeed(AUTO_DRIVE_SPEED);
        } else {
          /*
              the tower is within range of the flame (checked by the tower servo position) and turns off the drive engines.
              it then says the starts extinguishing the flame
          */
          flameFound = 1;
          leftDrive->setSpeed(0);
          rightDrive->setSpeed(0);
        }
      }
    }




  }
  //if the pilot flame is supposed to move, do so
  pilotFlame->tick();
  //make sure the pressure is within its bounds
  pressureSustainer->tick();
  //delay 10ms because arduino
  delay(10);

}

/* callback function for I2C communication */
void receiveData(int byteCount){
  while (Wire.available()) {
    //while theres bytes to read (not package)
    int number = Wire.read();
    //check the current state of the protocol
    switch (state)
    {
    case ST_IDLE:
      //check the first byte wether it matches a header or not
      if (number == 0x5C)
          state = ST_HEAD1_OK;
      break;
      //check second byte to see if it matches
    case ST_HEAD1_OK:
      if (number == 0xF3)
          state = ST_HEAD2_OK;
      else
          state = ST_IDLE;
      break;
    case ST_HEAD2_OK:
      //check wether the third header matches
      if (number == 0x0A)
          state = ST_HEADER_OK;
      else
          state = ST_IDLE;
      break;
    case ST_HEADER_OK:
      //read the 4th byte, flame coordinates in x direction
      cameraCoordinates[0] = number - 127;
      state = ST_X_READ;
      break;
    case ST_X_READ:
      //read the 5th byte, flame coordinates in y direction
      cameraCoordinates[1] = number - 127;
      //if both coordinates are not 127 (value if no flame is detected), new coordinates are received
      if(cameraCoordinates[0] < 127 && cameraCoordinates[1] < 127) newCoordinatesReceived = 1;
      //else no new coordinates
      else newCoordinatesReceived = 0;
      state = ST_IDLE;
      break;
    default:
      break;
    }
  }
}