#include "ugv.hh"

#include "ugv_comms.hh"
#include "ugv_display.hh"
#include "ugv_io.hh"

namespace ugv {

static const char *TAG = "ugv_main";

//#define CALIBRATE

static constexpr float PI =
    3.1415926535897932384626433832795028841971693993751058209749445923078164062;

constexpr uint64_t LOOP_PERIOD_US = 1e6 / 100;
constexpr float LOOP_HZ = 1000000.f / static_cast<float>(LOOP_PERIOD_US);
constexpr float LOOP_PERIOD_S = 1 / LOOP_HZ;
constexpr TickType_t WATCHDOG_TICKS = pdMS_TO_TICKS(20 * 1000);

constexpr uint64_t NOISE_PERIOD_US = 2000e3;
constexpr float NOISE_PERIOD_S = static_cast<float>(NOISE_PERIOD_US) * 1.e-6f;

constexpr float ACCEL_NOISE_THRESHOLD = 0.001;
constexpr float GYRO_NOISE_THRESHOLD = 0.3;

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);
}

extern "C" void UGV_TickTimeout(void *arg) {
  UGV *ugv = (UGV *)arg;
  ugv->OnTick();
}

UGV::UGV() : angle_controller_(LOOP_PERIOD_S) {
  SetTarget({34.069022, -118.443067});

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

  SetConfig(DefaultConfig());
}

config::Config UGV::DefaultConfig() {
  config::Config c;

  auto *apid = c.mutable_angle_pid();
  apid->set_kp(0.06);
  apid->set_ki(0.01);
  apid->set_kd(0.4);
  apid->set_max_output(0.6);
  apid->set_max_i_error(15.0);

  c.set_min_target_dist(1.0);
  c.set_min_flip_pitch(90.0);
  c.set_drive_power(0.4);
  c.set_mag_declination(11.5);
  return c;
}

void UGV::SetConfig(const config::Config &conf) {
  auto &apid = conf.angle_pid();
  angle_controller_.SetPID(apid.kp(), apid.ki(), apid.kd());
  angle_controller_.MaxOutput(apid.max_output());
  angle_controller_.MaxIError(apid.max_i_error());
  conf_ = conf;
}

void UGV::SetTarget(LatLong target) { target_ = target; }

void UGV::Init() {
  esp_timer_init();

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

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

  esp_timer_create_args_t timer_args;
  timer_args.callback = UGV_TickTimeout;
  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;
  last_noise_ = 0;
  last_left_ = 0;
  last_right_ = 0;
}

void UGV::UpdateAHRS() {
  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);
  yaw_ = ahrs_.getYaw() + conf_.mag_declination();
  while (yaw_ > 360.) {
    yaw_ -= 360.;
  }
  while (yaw_ < 0.) {
    yaw_ += 360.;
  }
  pitch_ = ahrs_.getPitch();
  roll_ = ahrs_.getRoll();
}

void UGV::DoDebugPrint() {
  auto &mpu = inputs_.mpu;
#ifdef CALIBRATE
  ESP_LOGD(TAG, "inputs: acc=(%f, %f, %f) gyro=(%f, %f, %f) mag=(%f, %f, %f)",
           mpu.accel.x, mpu.accel.y, mpu.accel.z, mpu.gyro_rate.x,
           mpu.gyro_rate.y, mpu.gyro_rate.z, mpu.mag.x, mpu.mag.y, mpu.mag.z);
#else
  ESP_LOGD(TAG, "inputs: acc=(%f, %f, %f) gyro=(%f, %f, %f) mag=(%f, %f, %f)",
           mpu.accel.x, mpu.accel.y, mpu.accel.z, mpu.gyro_rate.x,
           mpu.gyro_rate.y, mpu.gyro_rate.z, mpu.mag.x, mpu.mag.y, mpu.mag.z);
  ESP_LOGD(TAG, "ahrs: yaw=%f, pitch=%f, roll=%f", yaw_, pitch_, roll_);
  ESP_LOGV(TAG, "target angle: %f", angle_controller_.Setpoint());
#endif
}

void UGV::ReadComms() {
  comms->Lock();
  UpdateLocationFromGPS(*(comms->status.mutable_location()), inputs_.gps);
  comms->status.set_yaw_angle(yaw_);
  comms->status.set_pitch_angle(pitch_);
  comms->status.set_roll_angle(roll_);
  comms->status.set_is_still(is_still_);
  UGV_State comms_ugv_state = comms->status.state();
  TickType_t ticks_since_last_packet =
      (xTaskGetTickCount() - comms->last_packet_tick);
  if (ticks_since_last_packet >= WATCHDOG_TICKS &&
      (comms_ugv_state != UGV_State::STATE_IDLE ||
       current_state_ != UGV_State::STATE_IDLE)) {
    ESP_LOGW(TAG, "No comms for %f ticks, disabling",
             ticks_since_last_packet * (1000.f / pdMS_TO_TICKS(1000)));
    comms_ugv_state = UGV_State::STATE_IDLE;
  }
  if (comms_ugv_state != current_state_) {
    ESP_LOGI(TAG, "Comms updated UGV state to %d", comms_ugv_state);
    current_state_ = comms_ugv_state;
  }
  if (comms->new_target) {
    SetTarget(*comms->new_target);
    ESP_LOGI(TAG, "Updating target to (%f, %f)", target_.latitude,
             target_.longitude);
    delete comms->new_target;
    comms->new_target = nullptr;
  }
  if (comms->new_config) {
    ESP_LOGI(TAG, "Updating config");
    SetConfig(*comms->new_config);
    delete comms->new_config;
    comms->new_config = nullptr;
  }
  comms->Unlock();
}

void UGV::OnTick() {
  ESP_LOGV(TAG, "OnTick");
  int64_t time_us = esp_timer_get_time();
  float time_s = ((float)time_us) / 1e6;
  io->ReadInputs(inputs_);
  UpdateAHRS();

  angle_controller_.Input(yaw_);
  float drive_power = 0.;
  outputs_.left_motor = 0.0;
  outputs_.right_motor = 0.0;

  float pitch = roll_;  // TODO: fix quaternion -> YPR
  auto min_flip_pitch = conf_.min_flip_pitch();
  bool is_upside_down = (pitch > min_flip_pitch) || (pitch < -min_flip_pitch);

  io::Vec3f d_accel = inputs_.mpu.accel - last_mpu_.accel;
  io::Vec3f d_gyro = inputs_.mpu.gyro_rate - last_mpu_.gyro_rate;
  last_mpu_ = inputs_.mpu;

  accel_noise_accum_ += d_accel.norm2() * LOOP_PERIOD_S;
  gyro_noise_accum_ += d_gyro.norm2() * LOOP_PERIOD_S;

  if (time_us >= last_noise_ + NOISE_PERIOD_US) {
    accel_noise_ = accel_noise_accum_ / NOISE_PERIOD_S;
    gyro_noise_ = gyro_noise_accum_ / NOISE_PERIOD_S;
    is_still_ = (accel_noise_ < ACCEL_NOISE_THRESHOLD) &&
                (gyro_noise_ < GYRO_NOISE_THRESHOLD);
    ESP_LOGD(TAG, "is still: %s, accel noise: %f, gyro noise: %f",
             is_still_ ? "still" : "moving", accel_noise_, gyro_noise_);
    accel_noise_accum_ = 0;
    gyro_noise_accum_ = 0;
    last_noise_ = time_us;
  }

  ReadComms();
  next_state_ = current_state_;

  switch (current_state_) {
    default:
      ESP_LOGW(TAG, "unhandled state: %d", current_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 = 0.8;
      outputs_.right_motor = 0.8;
      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 <= conf_.min_target_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);
      drive_power = conf_.drive_power();
      break;
    }
    case UGV_State::STATE_TEST:
#define BASIC_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;
  }

  constexpr float MAX_ACCEL = 16 * LOOP_PERIOD_S;

  if (inputs_.mpu.gyro_rate.x == 0 && inputs_.mpu.gyro_rate.y == 0 &&
      inputs_.mpu.gyro_rate.z == 0) {
    outputs_.left_motor = 0;
    outputs_.right_motor = 0;
    did_miss_mpu_ = true;
  } else {
    float dleft = outputs_.left_motor - last_left_;
    float dright = outputs_.right_motor - last_right_;
    if (dleft > MAX_ACCEL) {
      outputs_.left_motor = last_left_ + MAX_ACCEL;
    } else if (dleft < -MAX_ACCEL) {
      outputs_.left_motor = last_left_ - MAX_ACCEL;
    }
    if (dright > MAX_ACCEL) {
      outputs_.right_motor = last_right_ + MAX_ACCEL;
    } else if (dright < -MAX_ACCEL) {
      outputs_.right_motor = last_right_ - MAX_ACCEL;
    }
  }
  io->WriteOutputs(outputs_);
  last_left_ = outputs_.left_motor;
  last_right_ = outputs_.right_motor;

  if (current_state_ != next_state_) {
    ESP_LOGI(TAG, "switching state to %d", next_state_);
  }
  comms->Lock();
  comms->status.set_state(next_state_);
  comms->Unlock();

#ifdef CALIBRATE
  const long debug_print_period = 50 * 1000;
#else
  const long debug_print_period = 500 * 1000;
#endif

  if (time_us >= last_print_ + debug_print_period) {
    if (did_miss_mpu_) {
      ESP_LOGW(TAG, "no MPU data, disabling");
      did_miss_mpu_ = false;
    }
    DoDebugPrint();
    last_print_ = time_us;
  }
}

UGV *the_ugv;

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

}  // namespace ugv