#include <Arduino.h>
#include <Wire.h>
#include <HardwareSerial.h>
#include <RunningMedian.h>

#define LED_PIN 2 // should be using builtin LED
#define LED_INTERVAL_OFF 2500 // in ms
#define LED_INTERVAL_ON 500 // in ms

#define SRV_PWM_FREQ 50
#define SRV_PWM_RES_BITS 10
#define SRV_PWM_REFRESH_USEC 20000

#define PAN_SRV_CHAN 1
#define PAN_SRV_PIN 32
#define PAN_SRV_PWM_MIN 1000
#define PAN_SRV_PWM_DEFAULT 1505 // the trim is on 77% duty
#define PAN_SRV_PWM_MAX 2000
#define PAN_SRV_PID_ITVL 100 // controller compute interval

#define TILT_SRV_CHAN 2
#define TILT_SRV_PIN 33
#define TILT_SRV_PWM_MIN 1180 // corresp with 0 deg
#define TILT_SRV_PWM_DEFAULT TILT_SRV_PWM_MIN
#define TILT_SRV_PWM_MAX 2090 // corresp with 90 deg

#define MAG_ADDRESS 0x1E
#define MAG_REG_CONFIG_A 0x00
#define MAG_REG_CONFIG_B 0x01
#define MAG_REG_MODE 0x02
#define MAG_REG_DATA 0x03
#define MAG_REG_IDTF 0x0A
#define MAG_PRINT_NUM 300 // approximately 30 seconds

#define GPS_BAUDRATE 9600
#define GPS_RX_PIN 19 // connected to GPS' TX pin
#define GPS_TX_PIN 18 // connected to GPS' RX pin (optional)

#define CONFIG_FREQ_HB 3
#define CONFIG_FREQ_AGL 4
#define CONFIG_MAG_UPD_FREQ 7
#define CONFIG_USE_GPS 8
#define CONFIG_KP 11
#define CONFIG_KI 12
#define CONFIG_KD 13

#define EARTH_RADIUS 6371

#define SERIAL_BAUDRATE 115200
#define CMD_PARAMS_NUM 3

enum Mode {
  OFF, // antenna tracker movement off
  LOCATION, // automatically determine the angles by UAV location
  SETPOINT_ANGLE, // manually set bearing & elevation angles
  SERVO_PWM // manually set pan & tilt servo pwm
};

struct Location {
  double latitude;
  double longitude;
  float altitude; // in meters
};

HardwareSerial SerialGPS(2);
RunningMedian bearingSamples(5); // remove spikes from data

uint32_t timePassed = millis();
bool isMagExists = false;
uint16_t magPrintRemaining = 0;

Mode mode = OFF;
Location atLoc = {
  -7.7738408, // location default to D13 (GAMAFORCE HQ)
  110.3783263,
  1.00 // altitude above ground level in meters
}; // antenna tracker location
Location uavLoc; // UAV/vehicle location
float bearActual = -1.0f; // actual bearing from sensor
float bearSetpoint; // bearing/azimuth angle setpoint
float elevSetpoint; // elevation angle setpoint
uint16_t panPwm = PAN_SRV_PWM_DEFAULT;
uint16_t tiltPwm = TILT_SRV_PWM_DEFAULT;
struct {
  uint8_t freqHB = 1; // in Hz
  uint8_t freqAGL = 10;
  float Kp = 3.1f;
  float Ki = 0.0f;
  float Kd = 0.02f;
  uint8_t magUpdFreq = 20; // magnetometer update freq
  float magDec = 1.46607f; // magnetic declination
  bool useGpsAsLoc = false; // override AT location using GPS reading
} configs;

int64_t strToInt64(String str) {
  uint8_t DIGITS_PER_PART = 9; // int32 max has >9 digits
  bool isNegative = str.charAt(0) == '-';
  if (isNegative) str = str.substring(1);
  uint8_t digitsLen = str.length();
  int64_t result = 0;
  
  for (int8_t e = digitsLen; e > 0; e -= DIGITS_PER_PART) {
    int8_t pos = (digitsLen - e) / DIGITS_PER_PART;
    int8_t s = e - DIGITS_PER_PART;
    if (s < 0) s = 0;

    int64_t partNum = str.substring(s, e).toInt();
    for (int8_t p = 0; p < pos; p++) partNum *= 1000000000;

    result += partNum;
  }

  if (isNegative) result *= -1;

  return result;
}

uint16_t usPwmToTicks(uint16_t usec) {
  return (uint16_t)((float)usec *
    (pow(2, SRV_PWM_RES_BITS) / (float)SRV_PWM_REFRESH_USEC));
}

// convert lat/lon format from DMM to DD
double convertDmmStrToDd (String &dmmStr) {
  double dmm = dmmStr.toDouble();
  int16_t numPart = static_cast<uint16_t>(dmm) / 100;
  double fracPart = dmm - static_cast<double>(numPart * 100);
  fracPart = fracPart / 60.0;
  double dd = static_cast<double>(numPart) + fracPart;

  if (dmmStr.endsWith("S")) dd *= -1.0; // for latitude
  if (dmmStr.endsWith("W")) dd *= -1.0; // for longitude

  return dd;
}

void sendMsgHeartbeat() {
  Serial.print("HB,");
  Serial.print(timePassed);
  Serial.print(",");
  Serial.println(mode);
}

void sendMsgAngles() {
  Serial.print("AGL,");
  Serial.print(bearActual);
  Serial.print(",");
  Serial.print(bearSetpoint);
  Serial.print(",");
  Serial.println(elevSetpoint);
}

void sendMsgUAVLocation() {
  Serial.print("UAV,");
  Serial.print(uavLoc.latitude, 7);
  Serial.print(",");
  Serial.print(uavLoc.longitude, 7);
  Serial.print(",");
  Serial.println(uavLoc.altitude, 3);
}

void sendMsgATLocation() {
  Serial.print("AT,");
  Serial.print(atLoc.latitude, 7);
  Serial.print(",");
  Serial.print(atLoc.longitude, 7);
  Serial.print(",");
  Serial.println(atLoc.altitude, 3);
}

void sendMsgSetpointAngles() {
  Serial.print("SP,");
  Serial.print(bearSetpoint);
  Serial.print(",");
  Serial.println(elevSetpoint);
}

void sendMsgServoPwm() {
  Serial.print("SRV,");
  Serial.print(panPwm);
  Serial.print(",");
  Serial.println(tiltPwm);
}

void sendMsgPIDGains() {
  Serial.print("PID,");
  Serial.print(configs.Kp, 3);
  Serial.print(",");
  Serial.print(configs.Ki, 3);
  Serial.print(",");
  Serial.println(configs.Kd, 3);
}

void sendMsgGPSReading(double latitude, double longitude,
  float altitude, float hdop) {
  Serial.print("GPS,");
  Serial.print(latitude, 7);
  Serial.print(",");
  Serial.print(longitude, 7);
  Serial.print(",");
  Serial.print(altitude, 2);
  Serial.print(",");
  Serial.println(hdop, 2);
}

void setPanServoPwm(uint16_t pwmValue) {
  panPwm = pwmValue;
  #ifdef ESP32
  ledcWrite(PAN_SRV_CHAN, usPwmToTicks(pwmValue));
  #endif
}

void setTiltServoPwm(uint16_t pwmValue) {
  tiltPwm = pwmValue;
  #ifdef ESP32
  ledcWrite(TILT_SRV_CHAN, usPwmToTicks(pwmValue));
  #endif
}

void resetServosPosition() {
  setPanServoPwm(PAN_SRV_PWM_DEFAULT);
  setTiltServoPwm(TILT_SRV_PWM_DEFAULT);
}

void initServo() {  
  #ifdef ESP32
  // setup pan servo
  ledcSetup(PAN_SRV_CHAN, SRV_PWM_FREQ, SRV_PWM_RES_BITS);
  ledcAttachPin(PAN_SRV_PIN, PAN_SRV_CHAN);

  // setup tilt servo
  ledcSetup(TILT_SRV_CHAN, SRV_PWM_FREQ, SRV_PWM_RES_BITS);
  ledcAttachPin(TILT_SRV_PIN, TILT_SRV_CHAN);
  #endif

  resetServosPosition();
}

void setBearingAngleSetpoint(float deg) {
  bearSetpoint = deg;
}

void setElevationAngleSetpoint(float deg) {
  elevSetpoint = deg;
  
  // TODO: implement algorithm to smoothing servo movement
  uint16_t newPwm = map(deg, 0, 90, TILT_SRV_PWM_MIN, TILT_SRV_PWM_MAX);
  setTiltServoPwm(newPwm);
}

void calculateNewSetpointAngles() {
  double startLat = radians(atLoc.latitude);
  double startLng = radians(atLoc.longitude);
  double destLat = radians(uavLoc.latitude);
  double destLng = radians(uavLoc.longitude);

  double y = sin(destLng - startLng) * cos(destLat);
  double x = cos(startLat) * sin(destLat) -
    sin(startLat) * cos(destLat) * cos(destLng - startLng);
  float bear = degrees(atan2(y, x));
  bear = bear < 0 ? bear + 360.0 : bear;

  double dist = EARTH_RADIUS * acos(
    sin(startLat) * sin(destLat) +
    cos(startLat) * cos(destLat) * cos(destLng - startLng));
  float altDiff = uavLoc.altitude - atLoc.altitude;
  float elev = degrees(atan2(altDiff, dist * 1000.0));
  elev = elev < 0 ? 0 : elev;

  setBearingAngleSetpoint(bear);
  setElevationAngleSetpoint(elev);
}

void setMode(Mode newMode) {
  mode = newMode;
  if (newMode == OFF) {
    resetServosPosition();
  }
  if (newMode == LOCATION) {
    calculateNewSetpointAngles();
  }
}

void scanI2CDevices() {
  uint8_t nDevices = 0;

  Serial.println("> Scanning I2C devices");

  for (byte address = 1; address < 127; address++) {
    Wire.beginTransmission(address);
    byte error = Wire.endTransmission();

    if (error == 0) {
      Serial.print("> Found at address 0x");
      Serial.println(address, HEX);
      nDevices++;
    } else if (error == 4) {
      Serial.print("> Unknown error at address 0x");
      Serial.println(address, HEX);
    }
  }

  Serial.print("> I2C devices found: ");
  Serial.println(nDevices);
}

// blink LED every n seconds
void blinkLED() {
  static uint32_t prevTime;
  uint16_t blinkInterval = mode == OFF
    ? LED_INTERVAL_OFF
    : LED_INTERVAL_ON;

  if (timePassed - prevTime >= blinkInterval) {
    digitalWrite(LED_PIN, HIGH);
    prevTime = timePassed;
  }
  if (timePassed - prevTime >= 100) {
    digitalWrite(LED_PIN, LOW);
  }
}

void printHeartbeatMsg() {
  static uint32_t prevTime;

  if (configs.freqHB == 0) return;
  if (timePassed - prevTime >= (1000 / configs.freqHB)) {
    sendMsgHeartbeat();
    prevTime = timePassed;
  }
}

void printAnglesMsg() {
  static uint32_t prevTime;

  if (configs.freqAGL == 0 || mode == OFF) return;
  if (timePassed - prevTime >= (1000 / configs.freqAGL)) {
    sendMsgAngles();
    prevTime = timePassed;
  }
}

void initMagnetometer() {
  Wire.beginTransmission(MAG_ADDRESS);
  Wire.write(MAG_REG_IDTF);
  byte error = Wire.endTransmission();

  // check if communication with the I2C device success
  if (error != 0) return;

  Wire.requestFrom(MAG_ADDRESS, 3);
  // HMC5883l identification bytes
  if (Wire.read() != 'H') return;
  if (Wire.read() != '4') return;
  if (Wire.read() != '3') return;

  isMagExists = true;

  // configure mag/compass
  Wire.beginTransmission(MAG_ADDRESS);
  Wire.write(MAG_REG_CONFIG_A);
  Wire.write(0x70);
  Wire.endTransmission();
  
  Wire.beginTransmission(MAG_ADDRESS);
  Wire.write(MAG_REG_CONFIG_B);
  Wire.write(0xA0);
  Wire.endTransmission();
  
  Wire.beginTransmission(MAG_ADDRESS);
  Wire.write(MAG_REG_MODE);
  Wire.write(0x00);
  Wire.endTransmission();
}

void readMagnetometer() {
  static uint32_t prevTime;

  if (!isMagExists) return;
  if (timePassed - prevTime >= (1000 / configs.magUpdFreq)) {
    int16_t x, y, z;

    Wire.beginTransmission(MAG_ADDRESS);
    Wire.write(MAG_REG_DATA); // start to read from this register
    Wire.endTransmission();

    byte retNum = Wire.requestFrom(MAG_ADDRESS, 6);
    if (Wire.available() < 6) return;

    x = Wire.read() << 8; // MSB x
    x |= Wire.read(); // LSB x
    z = Wire.read() << 8; // MSB z
    z |= Wire.read(); // LSB z
    y = Wire.read() << 8; // MSB y
    y |= Wire.read(); // LSB y

    float bearing = atan2(y, x);
    bearing += configs.magDec; // account for declination angle
    if (bearing < 0) bearing += 2*PI; // correct reversed signs
    if (bearing > 2*PI) bearing -= 2*PI; // check for wrap
    bearing = 2*PI - bearing; // reverse rotation

    float bearingDeg = degrees(bearing);
    bearingSamples.add(bearingDeg);

    // Serial.print("Raw: ");
    // Serial.print(x); Serial.print(" ");
    // Serial.print(y); Serial.print(" ");   
    // Serial.print(z);
    // Serial.print("\tBear: ");
    // Serial.println(bearingDeg);

    if (magPrintRemaining > 0) {
      uint16_t printProgress = MAG_PRINT_NUM-magPrintRemaining+1;
      float percentProgress = (float)printProgress/MAG_PRINT_NUM*100;
      Serial.print("> ");
      Serial.print(x); Serial.print(",");
      Serial.print(y); Serial.print(",");
      Serial.print(z); Serial.print(",");
      Serial.print(percentProgress, 1);
      Serial.println("%");
      magPrintRemaining--;
    }

    bearActual = bearingSamples.getMedian();
    prevTime = timePassed;
  }
}

void processNmeaMessage(String &msg) {
  msg.setCharAt(2, 'X'); // for compatibility with other GNSS sats

  if (msg.startsWith("$GXGGA")) {
    struct {
      String latitude, longitude, heightMsl, geoidSep, hdop;
    } usefulParts;
    uint8_t msgFieldIdx = 0;

    for (uint16_t i = 0; i < msg.length(); i++) {
      char &c = msg[i];
      if (c == ',') {
        msgFieldIdx++;
        continue;
      }
      if (msgFieldIdx == 2 || msgFieldIdx == 3)
        usefulParts.latitude += c;
      if (msgFieldIdx == 4 || msgFieldIdx == 5)
        usefulParts.longitude += c;
      if (msgFieldIdx == 9)
        usefulParts.heightMsl += c;
      if (msgFieldIdx == 11)
        usefulParts.geoidSep += c;
      if (msgFieldIdx == 8)
        usefulParts.hdop += c;
    }

    float hdop = usefulParts.hdop.toFloat();

    if (hdop < 10.0 && usefulParts.latitude != "") {
      double latitude = convertDmmStrToDd(usefulParts.latitude);
      double longitude = convertDmmStrToDd(usefulParts.longitude);
      float heightMsl = usefulParts.heightMsl.toFloat();
      float geoidSep = usefulParts.geoidSep.toFloat();
      float altitude = heightMsl + geoidSep; // ellipsoidal height

      sendMsgGPSReading(latitude, longitude, altitude, hdop);

      // a good HDOP to use is considered to be below 2
      if (configs.useGpsAsLoc && hdop < 2.0) {
        atLoc.latitude = latitude;
        atLoc.longitude = longitude;
      }
    }
  }
}

void readGPS() {
  static bool isReceivingNmea = false;
  static String receivedNmea = "";

  if (SerialGPS.available() > 0) {
    char rc = SerialGPS.read();

    if (!isReceivingNmea && rc == '$') {
      isReceivingNmea = true;
    }

    if (isReceivingNmea) {
      if (rc != '\n') {
        receivedNmea += rc;
      } else {
        receivedNmea.trim();
        processNmeaMessage(receivedNmea);

        isReceivingNmea = false;
        receivedNmea = "";
      }
    }
  }
}

void runCommand(String &cmd) {
  uint16_t cmdStrLen = cmd.length();
  char cmdType = '\0';
  int64_t cmdParams[CMD_PARAMS_NUM];
  String paramStrTemp = "";
  uint8_t paramIdx = 0;

  // fill array with the same initial value
  for(uint8_t i = 0; i < CMD_PARAMS_NUM; i++) {
    cmdParams[i] = INT64_MIN;
  }

  for (uint16_t i = 0; i < cmdStrLen && paramIdx < CMD_PARAMS_NUM; i++) {
    char c = cmd[i];
    if (i == 0) {
      cmdType = c;
      continue;
    }
    if (i == 1 && c == ',') continue;
    if ((c >= '0' && c <= '9') || c == '-') {
      paramStrTemp += c;
    }
    if (c == ',' || i == cmdStrLen-1) {
      cmdParams[paramIdx] = strToInt64(paramStrTemp);
      paramStrTemp = "";
      paramIdx++;
    }
  }

  int64_t &p1 = cmdParams[0];
  int64_t &p2 = cmdParams[1];
  int64_t &p3 = cmdParams[2];

  switch(cmdType) {
    case 'M': // set mode
      if (p1 >= 0 && p1 <= 3) {
        setMode(static_cast<Mode>(p1));
      }
      break;
    case 'A': // set AT location
      if (p1 != INT64_MIN && p2 != INT64_MIN) {
        atLoc.latitude = static_cast<double>(p1) / pow(10, 7);
        atLoc.longitude = static_cast<double>(p2) / pow(10, 7);
        if (p3 != INT64_MIN) {
          atLoc.altitude = static_cast<float>(p3) / pow(10, 3);
        }
        if (mode == LOCATION) {
          calculateNewSetpointAngles();
        }
        sendMsgATLocation();
      }
      break;
    case 'V': // set vehicle location (mode == LOCATION)
      if (mode != LOCATION) break;
      if (p1 != INT64_MIN && p2 != INT64_MIN) {
        uavLoc.latitude = static_cast<double>(p1) / pow(10, 7);
        uavLoc.longitude = static_cast<double>(p2) / pow(10, 7);
        if (p3 != INT64_MIN) {
          uavLoc.altitude = static_cast<float>(p3) / pow(10, 3);
        }
        calculateNewSetpointAngles();
        sendMsgUAVLocation();
      }
      break;
    case 'B': // set bearing setpoint angle (mode == SETPOINT_ANGLE)
      if (mode != SETPOINT_ANGLE) break;
      if (p1 >= 0 && p1 <= 359) {
        setBearingAngleSetpoint(static_cast<float>(p1));
        sendMsgSetpointAngles();
      }
      break;
    case 'E': // set elevation setpoint angle (mode == SETPOINT_ANGLE)
      if (mode != SETPOINT_ANGLE) break;
      if (p1 >= 0 && p1 <= 90) {
        setElevationAngleSetpoint(static_cast<float>(p1));
        sendMsgSetpointAngles();
      }
      break;
    case 'P': // set servo pan pwm value (mode == SERVO_PWM)
      if (mode != SERVO_PWM) break;
      if (p1 >= PAN_SRV_PWM_MIN && p1 <= PAN_SRV_PWM_MAX) {
        setPanServoPwm(p1);
        sendMsgServoPwm();
      }
      break;
    case 'T': // set servo tilt pwm value (mode == SERVO_PWM)
      if (mode != SERVO_PWM) break;
      if (p1 >= TILT_SRV_PWM_MIN && p1 <= TILT_SRV_PWM_MAX) {
        setTiltServoPwm(p1);
        sendMsgServoPwm();
      }
      break;
    case 'C': // set config
      if (p1 == INT64_MIN || p2 == INT64_MIN) break;
      if (p1 == CONFIG_FREQ_HB) {
        configs.freqHB = p2;
      }
      if (p1 == CONFIG_FREQ_AGL) {
        configs.freqAGL = p2;
      }
      if (p1 == CONFIG_MAG_UPD_FREQ) {
        configs.magUpdFreq = p2;
      }
      if (p1 == CONFIG_USE_GPS) {
        configs.useGpsAsLoc = static_cast<bool>(p2);
      }
      if (p1 >= CONFIG_KP && p1 <= CONFIG_KD) { // PID gains
        float newGain = static_cast<float>(p2) / pow(10, 3);
        if (p1 == CONFIG_KP) configs.Kp = newGain;
        if (p1 == CONFIG_KI) configs.Ki = newGain;
        if (p1 == CONFIG_KD) configs.Kd = newGain;
        sendMsgPIDGains();
      }
      break;
    case 'X': // run specific command
      if (p1 == 0) { // request all messages
        sendMsgHeartbeat();
        sendMsgAngles();
        sendMsgUAVLocation();
        sendMsgATLocation();
        sendMsgSetpointAngles();
        sendMsgServoPwm();
        sendMsgPIDGains();
      }
      if (p1 == 1) { // scan I2C address
        scanI2CDevices();
      }
      if (p1 == 2) { // print raw mag readings
        if (magPrintRemaining == 0) {
          magPrintRemaining = MAG_PRINT_NUM;
        } else {
          // stop printing
          magPrintRemaining = 0;
        }
      }
      break;
  }
}

// read command from serial
void readSerial() {
  static bool isReceivingCmd = false;
  static String receivedCmd = "";

  if (Serial.available() > 0) {
    char rc = Serial.read();

    if (!isReceivingCmd && rc >= 'A' && rc <= 'Z') {
      isReceivingCmd = true;
    }

    if (isReceivingCmd) {
      if (rc != '\n') {
        receivedCmd += rc;
      } else {
        runCommand(receivedCmd);

        isReceivingCmd = false;
        receivedCmd = "";
      }
    }
  }
}

void controlPanPID() {
  static uint32_t prevTime;
  static float error, integralError, deltaError, prevError;

  if (mode != LOCATION && mode != SETPOINT_ANGLE) return;
  if (bearActual == -1.0f) return; // not getting actual mag data
  if (timePassed - prevTime >= PAN_SRV_PID_ITVL) {
    error = bearSetpoint - bearActual;
    if (error > 180) error = -1 * (360.0 - error);
    else if (error < -180) error = 360.0 + error;
    integralError += ((error + prevError) *
      ((float)PAN_SRV_PID_ITVL / 1000.0)) / 2;
    deltaError = abs(error - prevError);

    float proportional = configs.Kp * error;
    float integral = configs.Ki * integralError;
    float derivative = configs.Kd * deltaError;

    uint16_t outputPwm = PAN_SRV_PWM_DEFAULT;
    outputPwm += proportional + integral + derivative;
    outputPwm = constrain(outputPwm, PAN_SRV_PWM_MIN, PAN_SRV_PWM_MAX);

    setPanServoPwm(outputPwm);

    // Serial.print("actu: ");
    // Serial.print(bearActual);
    // Serial.print("\tsetp: ");
    // Serial.print(bearSetpoint);
    // Serial.print("\tout: ");
    // Serial.println(outputPwm);

    prevError = error;
    prevTime = timePassed;
  }
}

void setup() {
  pinMode(LED_PIN, OUTPUT);
  Serial.begin(SERIAL_BAUDRATE);
  #ifdef ESP32
  SerialGPS.begin(GPS_BAUDRATE, SERIAL_8N1, GPS_RX_PIN, GPS_TX_PIN);
  #endif
  Wire.begin();

  initMagnetometer();
  initServo();
}

void loop() {
  timePassed = millis();
  
  // don't run blocking codes here
  blinkLED();
  printHeartbeatMsg();
  printAnglesMsg();
  readMagnetometer();
  readGPS();
  readSerial();
  controlPanPID();

  delay(1); // prevent high processor usage
}