#include <esp_log.h>
#include <esp_timer.h>
#include "MadgwickAHRS.h"
#include "i2c_mutex.h"
#include "ugv_comms.hh"
#include "ugv_display.hh"
#include "ugv_io.hh"

#include <math.h>

namespace ugv {

using ugv::comms::CommsClass;
using ugv::comms::messages::UGV_State;
using ugv::io::IOClass;

static const char *TAG = "ugv_main";

extern "C" {
SemaphoreHandle_t i2c_mutex;
}

constexpr uint64_t LOOP_PERIOD_US = 1e6 / 100;
constexpr float LOOP_PERIOD_S = 1000000.f / static_cast<float>(LOOP_PERIOD_US);
static const float PI =
    3.1415926535897932384626433832795028841971693993751058209749445923078164062;

static const float RAD_PER_DEG = PI / 180.f;
// Radius of earth in meters
static const float EARTH_RAD = 6372795.f;

static const float DRIVE_POWER       = 0.5;
static const float ANGLE_P           = 0.02;
static const float ANGLE_I           = 0.000005;
static const float ANGLE_D           = -0.01;
static const float MAX_ANGLE_POWER   = 0.3;
static const float MAX_ANGLE_I_ERROR = 15.0;
static const float MIN_DIST          = 10.0;

extern "C" void OnTimeout(void *arg);

void UpdateLocationFromGPS(comms::messages::Location &location,
                           const io::GpsData &        gps_data) {
  location.set_fix_quality(gps_data.fix_quality);
  location.set_latitude(gps_data.latitude);
  location.set_longitude(gps_data.longitude);
  location.set_altitude(gps_data.altitude);
}

struct LatLong {
 public:
  float latitude;
  float longitude;

  inline LatLong(double latitude_, double longitude_)
      : latitude(latitude_), longitude(longitude_) {}

  /**
   * Return distance from this LatLong to target, in meters
   */
  float distance_to(const LatLong &target) const {
    float lat1  = latitude * RAD_PER_DEG;
    float lat2  = target.latitude * RAD_PER_DEG;
    float long1 = longitude * RAD_PER_DEG;
    float long2 = target.longitude * RAD_PER_DEG;
    float clat1 = cosf(lat1);
    float clat2 = cosf(lat2);
    float a     = powf(sinf((long2 - long1) / 2.f), 2.f) * clat1 * clat2 +
              powf(sinf((lat2 - lat1) / 2.f), 2.f);
    float d_over_r = 2 * atan2f(sqrtf(a), sqrtf(1 - a));
    return d_over_r * EARTH_RAD;
  }

  float bearing_toward(const LatLong &target) const {
    float dlong  = (target.longitude - longitude) * RAD_PER_DEG;
    float sdlong = sinf(dlong);
    float cdlong = cosf(dlong);
    float lat1   = latitude * RAD_PER_DEG;
    float lat2   = target.latitude * RAD_PER_DEG;
    float slat1  = sinf(lat1);
    float clat1  = cosf(lat1);
    float slat2  = sinf(lat2);
    float clat2  = cosf(lat2);
    float num    = sdlong * clat2;
    float denom  = (clat1 * slat2) - (slat1 * clat2 * cdlong);
    float course = atan2f(num, denom);
    if (course < 0.0) {
      course += 2 * PI;
    }
    return course / RAD_PER_DEG;
  }
};

class PIDController {
 public:
  explicit PIDController(float dt, float kp, float ki = 0., float kd = 0.);

  void  MaxOutput(float max_output) { max_output_ = max_output; }
  float MaxOutput() const { return max_output_; }

  void  MaxIError(float max_i_error) { max_i_error_ = max_i_error; }
  float MaxIError() const { return max_i_error_; }

  void  Setpoint(float setpoint) { setpoint_ = setpoint; }
  float Setpoint() const { return setpoint_; }

  void  Input(float input) { input_ = input; }
  float Input() const { return input_; };

  float Error() const {
    float error = setpoint_ - input_;
    // TODO: have this be configurable
    while (error < 180.f) error += 360.f;
    while (error > 180.f) error -= 360.f;
    return error;
  }
  float Output() const { return output_; };

  float Update();
  float Update(float input);
  float Update(float input, float setpoint);

  void Reset();
  void Enable(bool enable = true) { enabled_ = enable; }
  void Disable() { Enable(false); }
  bool Enabled() const { return enabled_; }

 private:
  static float clamp_mag(float x, float mag);

  float dt_;
  float kp_;
  float ki_;
  float kd_;
  float max_output_;
  float max_i_error_;

  bool  enabled_;
  float setpoint_;
  float input_;
  float output_;

  float integral_;
  float last_error_;
};

float PIDController::clamp_mag(float x, float max_mag) {
  if (x > max_mag)
    return max_mag;
  else if (x < -max_mag)
    return -max_mag;
  else
    return x;
}

PIDController::PIDController(float dt, float kp, float ki, float kd)
    : dt_(dt),
      kp_(kp),
      ki_(ki),
      kd_(kd),
      max_output_(INFINITY),
      max_i_error_(INFINITY),
      enabled_(false),
      setpoint_(0),
      input_(0),
      output_(0),
      integral_(0),
      last_error_(0) {}

void PIDController::Reset() {
  enabled_    = false;
  setpoint_   = 0.;
  input_      = 0.;
  output_     = 0.;
  integral_   = 0.;
  last_error_ = NAN;
}

float PIDController::Update() {
  output_ = 0.;
  if (!enabled_) {
    return output_;
  }
  float error = Error();

  if (fabsf(error) > max_i_error_) {
    integral_ = 0.;
  } else {
    integral_ += error * dt_;
  }

  output_ += kp_ * error;
  output_ += ki_ * integral_;
  if (!isnan(last_error_)) {
    output_ += kd_ * (error - last_error_);
  }
  output_ = clamp_mag(output_, max_output_);

  last_error_ = error;
  return output_;
}

float PIDController::Update(float input) {
  Input(input);
  return Update();
}

float PIDController::Update(float input, float setpoint) {
  Setpoint(setpoint);
  return Update(input);
}

struct State {
 public:
  CommsClass *       comms;
  IOClass *          io;
  DisplayClass *     display;
  esp_timer_handle_t timer_handle;

  io::Inputs    inputs;
  io::Outputs   outputs;
  int64_t       last_print;
  Madgwick      ahrs_;
  LatLong       target;
  PIDController angle_controller_;

  State()
      : target{34.069022, -118.443067},
        angle_controller_(LOOP_PERIOD_S, ANGLE_P, ANGLE_I, ANGLE_D) {
    angle_controller_.MaxOutput(MAX_ANGLE_POWER);
    angle_controller_.MaxIError(MAX_ANGLE_I_ERROR);

    comms   = new CommsClass();
    io      = new IOClass();
    display = new DisplayClass(comms);
  }

  void Init() {
    esp_timer_init();
    i2c_mutex = xSemaphoreCreateMutex();

    ahrs_.begin(LOOP_PERIOD_S);  // rough sample frequency

    io->Init();
    comms->Init();
    display->Init();

    esp_timer_create_args_t timer_args;
    timer_args.callback        = OnTimeout;
    timer_args.arg             = this;
    timer_args.dispatch_method = ESP_TIMER_TASK;
    timer_args.name            = "ugv_main_loop";
    esp_timer_create(&timer_args, &this->timer_handle);
    esp_timer_start_periodic(timer_handle, LOOP_PERIOD_US);
    last_print = 0;
  }

  void OnTick() {
    ESP_LOGV(TAG, "OnTick");
    int64_t time_us = esp_timer_get_time();
    float   time_s  = ((float)time_us) / 1e6;
    io->ReadInputs(inputs);
    {
      io::Vec3f &g = inputs.mpu.gyro_rate, &a = inputs.mpu.accel,
                &m = inputs.mpu.mag;
      ahrs_.update(g.x, g.y, g.z, a.x, a.y, a.z, m.x, m.y, m.z);
    }
    if (time_us >= last_print + 500 * 1000) {  // 1s
      ESP_LOGD(TAG,
               "inputs: acc=(%f, %f, %f) gyro=(%f, %f, %f) mag=(%f, %f, %f)",
               inputs.mpu.accel.x, inputs.mpu.accel.y, inputs.mpu.accel.z,
               inputs.mpu.gyro_rate.x, inputs.mpu.gyro_rate.y,
               inputs.mpu.gyro_rate.z, inputs.mpu.mag.x, inputs.mpu.mag.y,
               inputs.mpu.mag.z);
      ESP_LOGD(TAG, "ahrs: yaw=%f, pitch=%f, roll=%f", ahrs_.getYaw(),
               ahrs_.getPitch(), ahrs_.getRoll());
      ESP_LOGD(TAG, "PID: error: %f", angle_controller_.Error());
      last_print = time_us;
    }

    comms->Lock();
    UpdateLocationFromGPS(comms->location, inputs.gps);
    comms->yaw_angle    = ahrs_.getYaw();
    UGV_State ugv_state = comms->ugv_state;
    comms->Unlock();
    UGV_State next_state = ugv_state;

    angle_controller_.Input(ahrs_.getYaw());
    float drive_power   = 0.;
    outputs.left_motor  = 0.0;
    outputs.right_motor = 0.0;

    float pitch = ahrs_.getPitch();

    bool is_upside_down = (pitch > 90.) || (pitch < -90.);

    switch (ugv_state) {
      default:
        ESP_LOGW(TAG, "unhandled state: %d", ugv_state);
        // fall through
      case UGV_State::STATE_IDLE:
      case UGV_State::STATE_FINISHED: angle_controller_.Disable(); break;
      case UGV_State::STATE_AQUIRING: {
        if (is_upside_down) {
          next_state = UGV_State::STATE_FLIPPING;
          break;
        }
        angle_controller_.Disable();
        TickType_t current_tick    = xTaskGetTickCount();
        TickType_t ticks_since_gps = current_tick - inputs.gps.last_update;
        bool       not_old         = ticks_since_gps <= pdMS_TO_TICKS(2000);
        bool       not_invalid = inputs.gps.fix_quality != io::GPS_FIX_INVALID;
        if (not_old && not_invalid) {
          next_state = UGV_State::STATE_TURNING;
        }
        break;
      }
      case UGV_State::STATE_FLIPPING: {
        angle_controller_.Disable();
        outputs.left_motor = -1.0;
        outputs.right_motor = -1.0;
        if (!is_upside_down) {
          next_state = UGV_State::STATE_AQUIRING;
          break;
        }
        break;
      }
      case UGV_State::STATE_TURNING: {
        if (is_upside_down) {
          next_state = UGV_State::STATE_FLIPPING;
          break;
        }
        if (inputs.gps.fix_quality == io::GPS_FIX_INVALID) {
          next_state = UGV_State::STATE_AQUIRING;
          break;
        }

        LatLong current_pos = {inputs.gps.latitude, inputs.gps.longitude};
        float   tgt_bearing = current_pos.bearing_toward(target);
        angle_controller_.Enable();
        angle_controller_.Setpoint(tgt_bearing);

        if (fabs(angle_controller_.Error()) <= 5.0) {
          next_state = UGV_State::STATE_DRIVING;
        }
        break;
      }
      case UGV_State::STATE_DRIVING: {
        if (is_upside_down) {
          next_state = UGV_State::STATE_FLIPPING;
          break;
        }
        if (inputs.gps.fix_quality == io::GPS_FIX_INVALID) {
          next_state = UGV_State::STATE_AQUIRING;
          break;
        }

        LatLong current_pos = {inputs.gps.latitude, inputs.gps.longitude};
        float   tgt_dist    = current_pos.distance_to(target);

        if (tgt_dist <= MIN_DIST) {
          ESP_LOGI(TAG, "Finished driving to target");
          next_state = UGV_State::STATE_FINISHED;
          break;
        }

        float tgt_bearing = current_pos.bearing_toward(target);
        angle_controller_.Enable();
        angle_controller_.Setpoint(tgt_bearing);
        break;
      }
      case UGV_State::STATE_TEST:
#ifdef BASIC_TEST
        outputs.left_motor  = sinf(time_s * PI);
        outputs.right_motor = cosf(time_s * PI);
#else
        angle_controller_.Enable();
        angle_controller_.Setpoint(90.0);
#endif
        break;
      case UGV_State::STATE_DRIVE_HEADING:
        angle_controller_.Enable();
        angle_controller_.Setpoint(comms->drive_heading.heading());
        drive_power = comms->drive_heading.power();
        break;
    }

    if (angle_controller_.Enabled()) {
      float angle_pwr     = angle_controller_.Update();
      outputs.left_motor  = drive_power - angle_pwr;
      outputs.right_motor = drive_power + angle_pwr;
    }

    io->WriteOutputs(outputs);

    comms->Lock();
    comms->ugv_state = next_state;
    comms->Unlock();
  }
};

extern "C" void OnTimeout(void *arg) {
  State *state = (State *)arg;
  state->OnTick();
}

State *state;

void Setup(void) {
  ESP_LOGI(TAG, "Starting UAS UGV");
  state = new State();
  state->Init();
  ESP_LOGI(TAG, "Setup finished");
}

}  // namespace ugv

extern "C" void app_main() { ugv::Setup(); }