//
// Created by lgv on 2026/3/9.
//

#include "../include/s_curve_velocity_planner.h"

#include <cmath>
#include <limits>

namespace cmvr
{

SCurveVelocityPlanner1D::SCurveVelocityPlanner1D(double max_velocity,
                                                 double max_acceleration,
                                                 double max_jerk)
  : max_velocity_(std::abs(max_velocity))
  , max_acceleration_(std::abs(max_acceleration))
  , max_jerk_(std::abs(max_jerk))
  , target_replan_threshold_(DEFAULT_TARGET_REPLAN_THRESHOLD)
  , synchronize_replan_threshold_(DEFAULT_SYNCHRONIZE_REPLAN_THRESHOLD)
{
  reset();
}

void SCurveVelocityPlanner1D::setConstraints(double max_velocity,
                                             double max_acceleration,
                                             double max_jerk)
{
  max_velocity_ = std::abs(max_velocity);
  max_acceleration_ = std::abs(max_acceleration);
  max_jerk_ = std::abs(max_jerk);

  state_.velocity = clamp(state_.velocity, -max_velocity_, max_velocity_);
  state_.acceleration = clamp(state_.acceleration, -max_acceleration_, max_acceleration_);
  state_.target_velocity = clamp(state_.target_velocity, -max_velocity_, max_velocity_);

  if (state_.has_active_profile) {
    active_profile_ = calculateProfile(state_.velocity,
                                       state_.target_velocity,
                                       state_.acceleration);
    state_.elapsed_time = 0.0;
    state_.has_active_profile = active_profile_.valid && active_profile_.total_time > EPSILON;
  }

  updateIsMovingFlag();
}

void SCurveVelocityPlanner1D::getConstraints(double& max_velocity,
                                             double& max_acceleration,
                                             double& max_jerk) const
{
  max_velocity = max_velocity_;
  max_acceleration = max_acceleration_;
  max_jerk = max_jerk_;
}

void SCurveVelocityPlanner1D::setTargetReplanThreshold(double threshold)
{
  target_replan_threshold_ = std::max(0.0, threshold);
}

void SCurveVelocityPlanner1D::setSynchronizeReplanThreshold(double threshold)
{
  synchronize_replan_threshold_ = std::max(0.0, threshold);
}

void SCurveVelocityPlanner1D::initialize(double velocity, double acceleration)
{
  state_ = SCurveVelocityPlannerState();
  state_.velocity = clamp(velocity, -max_velocity_, max_velocity_);
  state_.acceleration = clamp(acceleration, -max_acceleration_, max_acceleration_);
  state_.jerk = 0.0;
  state_.target_velocity = state_.velocity;
  state_.elapsed_time = 0.0;
  state_.has_active_profile = false;

  active_profile_ = SCurveVelocityProfile();
  active_profile_.v0 = state_.velocity;
  active_profile_.a0 = state_.acceleration;
  active_profile_.vf = state_.velocity;
  active_profile_.valid = false;

  updateIsMovingFlag();
}

void SCurveVelocityPlanner1D::reset()
{
  state_ = SCurveVelocityPlannerState();
  active_profile_ = SCurveVelocityProfile();
}

void SCurveVelocityPlanner1D::overwriteState(double velocity,
                                             double acceleration,
                                             bool keep_target_velocity)
{
  state_.velocity = clamp(velocity, -max_velocity_, max_velocity_);
  state_.acceleration = clamp(acceleration, -max_acceleration_, max_acceleration_);
  state_.jerk = 0.0;
  state_.elapsed_time = 0.0;
  state_.has_active_profile = false;

  if (!keep_target_velocity) {
    state_.target_velocity = state_.velocity;
  } else {
    state_.target_velocity = clamp(state_.target_velocity, -max_velocity_, max_velocity_);
  }

  active_profile_ = SCurveVelocityProfile();
  active_profile_.v0 = state_.velocity;
  active_profile_.a0 = state_.acceleration;
  active_profile_.vf = state_.target_velocity;
  active_profile_.valid = false;

  updateIsMovingFlag();
}

void SCurveVelocityPlanner1D::synchronizeAndReplan(double velocity,
                                                   double acceleration)
{
  const double measured_velocity = clamp(velocity, -max_velocity_, max_velocity_);
  const double measured_acceleration =
      clamp(acceleration, -max_acceleration_, max_acceleration_);

  if (!state_.has_active_profile &&
      std::abs(measured_velocity - state_.target_velocity) <= VELOCITY_THRESHOLD) {
    state_.velocity = state_.target_velocity;
    state_.acceleration = 0.0;
    state_.jerk = 0.0;
    updateIsMovingFlag();
    return;
  }

  // If the measured state is still close to the active profile sample, keep
  // following that profile instead of restarting the plan every cycle.
  if (state_.has_active_profile &&
      std::abs(measured_velocity - state_.velocity) <= synchronize_replan_threshold_) {
    state_.velocity = measured_velocity;
    state_.acceleration = measured_acceleration;
    state_.jerk = getJerkAtTime(active_profile_, state_.elapsed_time);
    updateIsMovingFlag();
    return;
  }

  state_.velocity = measured_velocity;
  state_.acceleration = measured_acceleration;
  state_.jerk = 0.0;

  active_profile_ = calculateProfile(state_.velocity,
                                     state_.target_velocity,
                                     state_.acceleration);

  state_.elapsed_time = 0.0;
  state_.has_active_profile = active_profile_.total_time > EPSILON;

  updateIsMovingFlag();
}

void SCurveVelocityPlanner1D::setTargetVelocity(double target_velocity)
{
  const double new_target = clamp(target_velocity, -max_velocity_, max_velocity_);

  const bool same_target =
      std::abs(new_target - state_.target_velocity) <= target_replan_threshold_;

  if (same_target &&
      state_.has_active_profile) {
    return;
  }

  if (same_target &&
      !state_.has_active_profile &&
      std::abs(state_.velocity - new_target) <= VELOCITY_THRESHOLD &&
      std::abs(state_.acceleration) <= ACCELERATION_THRESHOLD) {
    return;
  }

  state_.target_velocity = new_target;

  active_profile_ = calculateProfile(state_.velocity,
                                     state_.target_velocity,
                                     state_.acceleration);

  state_.elapsed_time = 0.0;
  state_.has_active_profile = active_profile_.valid && active_profile_.total_time > EPSILON;

  updateIsMovingFlag();
}

double SCurveVelocityPlanner1D::getVelocityAtTime(const SCurveVelocityProfile& profile, double t) const
{
  if (!profile.valid) {
    return profile.v0;
  }

  if (t <= 0.0) {
    return profile.v0;
  }
  if (t >= profile.total_time) {
    return profile.vf;
  }

  double v = profile.v0;
  double a = profile.a0;

  auto process_segment = [&](double dt, double jerk) {
    v += a * dt + 0.5 * jerk * dt * dt;
    a += jerk * dt;
  };

  if (t <= profile.t1) {
    process_segment(t, profile.j1);
    return v;
  }
  process_segment(profile.t1, profile.j1);

  const double t_after_1 = t - profile.t1;
  if (t_after_1 <= profile.t2) {
    process_segment(t_after_1, profile.j2);
    return v;
  }
  process_segment(profile.t2, profile.j2);

  const double t_after_2 = t_after_1 - profile.t2;
  process_segment(t_after_2, profile.j3);
  return v;
}

double SCurveVelocityPlanner1D::getAccelerationAtTime(const SCurveVelocityProfile& profile, double t) const
{
  if (!profile.valid) {
    return profile.a0;
  }

  if (t <= 0.0) {
    return profile.a0;
  }
  if (t >= profile.total_time) {
    return 0.0;
  }

  double a = profile.a0;

  auto process_segment = [&](double dt, double jerk) {
    a += jerk * dt;
  };

  if (t <= profile.t1) {
    process_segment(t, profile.j1);
    return a;
  }
  process_segment(profile.t1, profile.j1);

  const double t_after_1 = t - profile.t1;
  if (t_after_1 <= profile.t2) {
    process_segment(t_after_1, profile.j2);
    return a;
  }
  process_segment(profile.t2, profile.j2);

  const double t_after_2 = t_after_1 - profile.t2;
  process_segment(t_after_2, profile.j3);
  return a;
}

double SCurveVelocityPlanner1D::getJerkAtTime(const SCurveVelocityProfile& profile, double t) const
{
  if (!profile.valid) {
    return 0.0;
  }

  if (t <= 0.0 || t >= profile.total_time) {
    return 0.0;
  }

  if (t <= profile.t1) {
    return profile.j1;
  }
  if (t <= profile.t1 + profile.t2) {
    return profile.j2;
  }
  return profile.j3;
}

void SCurveVelocityPlanner1D::integrateWholeProfile(const SCurveVelocityProfile& profile,
                                                    double& v_end,
                                                    double& a_end) const
{
  double v = profile.v0;
  double a = profile.a0;

  auto process_segment = [&](double dt, double jerk) {
    v += a * dt + 0.5 * jerk * dt * dt;
    a += jerk * dt;
  };

  process_segment(profile.t1, profile.j1);
  process_segment(profile.t2, profile.j2);
  process_segment(profile.t3, profile.j3);

  v_end = v;
  a_end = a;
}

bool SCurveVelocityPlanner1D::isVelocityLimitSatisfied(const SCurveVelocityProfile& profile) const
{
  if (!profile.valid) {
    return true;
  }

  std::vector<double> check_times;
  check_times.reserve(8);

  check_times.push_back(0.0);
  check_times.push_back(profile.t1);
  check_times.push_back(profile.t1 + profile.t2);
  check_times.push_back(profile.total_time);

  // If acceleration crosses zero inside segment 1, velocity may reach an
  // extremum within the segment and needs to be checked explicitly.
  if (std::abs(profile.j1) > EPSILON && profile.t1 > EPSILON) {
    const double tz = -profile.a0 / profile.j1;
    if (tz > EPSILON && tz < profile.t1 - EPSILON) {
      check_times.push_back(tz);
    }
  }

  for (double t : check_times) {
    const double v = getVelocityAtTime(profile, t);
    if (std::abs(v) > max_velocity_ + 1e-7) {
      return false;
    }
  }

  return true;
}

SCurveVelocityProfile SCurveVelocityPlanner1D::calculateProfile(double start_velocity,
                                                                double end_velocity,
                                                                double start_acceleration) const
{
  SCurveVelocityProfile best;
  best.v0 = clamp(start_velocity, -max_velocity_, max_velocity_);
  best.a0 = clamp(start_acceleration, -max_acceleration_, max_acceleration_);
  best.vf = clamp(end_velocity, -max_velocity_, max_velocity_);

  if (max_jerk_ <= EPSILON || max_acceleration_ <= EPSILON) {
    return best;
  }

  const double v0 = best.v0;
  const double a0 = best.a0;
  const double vf = best.vf;

  // Work in a signed frame where the requested delta velocity is non-negative.
  double s = sign(vf - v0);
  if (std::abs(vf - v0) <= VELOCITY_THRESHOLD) {
    s = 1.0;
  }

  const double dv = s * (vf - v0);
  const double a0s = s * a0;
  const double j = max_jerk_;
  const double amax = max_acceleration_;

  auto normalize_nonnegative = [](double x) {
    return (x < 0.0 && x > -1e-10) ? 0.0 : x;
  };

  auto try_candidate = [&](double t1, double t2, double t3,
                           double j1s, double j2s, double j3s) {
    t1 = normalize_nonnegative(t1);
    t2 = normalize_nonnegative(t2);
    t3 = normalize_nonnegative(t3);

    if (t1 < 0.0 || t2 < 0.0 || t3 < 0.0) {
      return;
    }

    SCurveVelocityProfile cand;
    cand.t1 = t1;
    cand.t2 = t2;
    cand.t3 = t3;
    cand.total_time = t1 + t2 + t3;

    cand.j1 = s * j1s;
    cand.j2 = s * j2s;
    cand.j3 = s * j3s;

    cand.v0 = v0;
    cand.a0 = a0;
    cand.vf = vf;
    cand.valid = true;

    double vend = 0.0;
    double aend = 0.0;
    integrateWholeProfile(cand, vend, aend);

    if (std::abs(vend - vf) > 1e-5) {
      return;
    }
    if (std::abs(aend) > 1e-5) {
      return;
    }
    if (!isVelocityLimitSatisfied(cand)) {
      return;
    }

    if (!best.valid || cand.total_time < best.total_time) {
      best = cand;
    }
  };

  {
    const double A2 = 0.5 * (2.0 * j * dv + a0s * a0s);
    if (A2 >= -1e-10) {
      const double A = std::sqrt(std::max(0.0, A2));
      if (A <= amax + 1e-10 && A >= a0s - 1e-10) {
        const double t1 = (A - a0s) / j;
        const double t2 = 0.0;
        const double t3 = A / j;
        try_candidate(t1, t2, t3, +j, 0.0, -j);
      }
    }
  }

  {
    const double A = amax;
    if (A > EPSILON && A >= a0s - 1e-10) {
      const double base = (2.0 * A * A - a0s * a0s) / (2.0 * j);
      const double t1 = (A - a0s) / j;
      const double t2 = (dv - base) / A;
      const double t3 = A / j;
      try_candidate(t1, t2, t3, +j, 0.0, -j);
    }
  }

  {
    const double B2 = 0.5 * (a0s * a0s - 2.0 * j * dv);
    if (B2 >= -1e-10) {
      const double B = std::sqrt(std::max(0.0, B2));
      if (B <= amax + 1e-10 && B >= -a0s - 1e-10) {
        const double t1 = (a0s + B) / j;
        const double t2 = 0.0;
        const double t3 = B / j;
        try_candidate(t1, t2, t3, -j, 0.0, +j);
      }
    }
  }

  {
    const double B = amax;
    if (B > EPSILON && B >= -a0s - 1e-10) {
      const double base = (a0s * a0s - 2.0 * B * B) / (2.0 * j);
      const double t1 = (a0s + B) / j;
      const double t2 = (base - dv) / B;
      const double t3 = B / j;
      try_candidate(t1, t2, t3, -j, 0.0, +j);
    }
  }

  if (!best.valid &&
      std::abs(vf - v0) <= VELOCITY_THRESHOLD &&
      std::abs(a0) <= ACCELERATION_THRESHOLD) {
    best.valid = false;
    best.total_time = 0.0;
    best.v0 = v0;
    best.a0 = a0;
    best.vf = vf;
  }

  return best;
}

void SCurveVelocityPlanner1D::generateTrajectory(const SCurveVelocityProfile& profile, double dt,
                                                 std::vector<double>& velocities,
                                                 std::vector<double>& accelerations,
                                                 std::vector<double>& jerks) const
{
  velocities.clear();
  accelerations.clear();
  jerks.clear();

  if (!profile.valid || profile.total_time <= 0.0 || dt <= 0.0) {
    velocities.push_back(profile.v0);
    accelerations.push_back(profile.a0);
    jerks.push_back(0.0);
    return;
  }

  const int num_points = static_cast<int>(std::ceil(profile.total_time / dt)) + 1;
  velocities.reserve(num_points);
  accelerations.reserve(num_points);
  jerks.reserve(num_points);

  for (double t = 0.0; t <= profile.total_time; t += dt) {
    velocities.push_back(getVelocityAtTime(profile, t));
    accelerations.push_back(getAccelerationAtTime(profile, t));
    jerks.push_back(getJerkAtTime(profile, t));
  }

  if (velocities.empty() || std::abs(velocities.back() - profile.vf) > VELOCITY_THRESHOLD) {
    velocities.push_back(profile.vf);
    accelerations.push_back(0.0);
    jerks.push_back(0.0);
  }
}

double SCurveVelocityPlanner1D::update(double dt)
{
  if (dt <= 0.0) {
    return state_.velocity;
  }

  if (!state_.has_active_profile || !active_profile_.valid) {
    state_.jerk = 0.0;
    updateIsMovingFlag();
    return state_.velocity;
  }

  state_.elapsed_time += dt;

  if (state_.elapsed_time >= active_profile_.total_time) {
    state_.elapsed_time = active_profile_.total_time;
    state_.velocity = active_profile_.vf;
    state_.acceleration = 0.0;
    state_.jerk = 0.0;
    state_.has_active_profile = false;
    updateIsMovingFlag();
    return state_.velocity;
  }

  state_.velocity = getVelocityAtTime(active_profile_, state_.elapsed_time);
  state_.acceleration = getAccelerationAtTime(active_profile_, state_.elapsed_time);
  state_.jerk = getJerkAtTime(active_profile_, state_.elapsed_time);

  updateIsMovingFlag();
  return state_.velocity;
}

void SCurveVelocityPlanner1D::updateIsMovingFlag()
{
  state_.is_moving =
      state_.has_active_profile ||
      std::abs(state_.velocity) > VELOCITY_THRESHOLD ||
      std::abs(state_.target_velocity) > VELOCITY_THRESHOLD ||
      std::abs(state_.acceleration) > ACCELERATION_THRESHOLD;
}

} // namespace cmvr
