#include <cmath>
#include "trapezoidal_trajectory.hpp"

TrapezoidalTrajectory::TrapezoidalTrajectory()
    : max_velocity_(1.0), max_acceleration_(0.25)
{
    max_accelerate_time_ = 4;
    max_accelerate_angle_ = max_velocity_ * max_accelerate_time_;

    nowTrajectoryIndex_ = -1;
    trajectoryFullFlag_ = 0;
    trajectoryHistory_[0].current_time_step = 0;
    min_angle_for_pose_stepping_max_speed_ = ARM_FOLLOWING_CHECKING_STEP * max_velocity_ * 2;
}

TrapezoidalTrajectory::~TrapezoidalTrajectory() {}

Trajectory* TrapezoidalTrajectory::addTrajectory()
{
    nowTrajectoryIndex_++;
    if (nowTrajectoryIndex_ >= MAX_TRAJECTORY_HISTORY) {
        nowTrajectoryIndex_ = 0;
        trajectoryFullFlag_ = 1;
    }
    return &trajectoryHistory_[nowTrajectoryIndex_];
}

void TrapezoidalTrajectory::computeMovingTrajectory(Trajectory *newTrajectory, const float *start, const float *end)
{
    float max_angle = 0.0;
    float angle_diff[USED_ARM_DOF];
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        newTrajectory->endAngles[i] = end[i];
        angle_diff[i] = end[i] - start[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
        }
    }
    if (max_angle < 0.01) {
        newTrajectory->accelerate_time_step = 0;
        newTrajectory->uniform_time_step = 0;
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            newTrajectory->aim_velocity[i] = 0;
            newTrajectory->now_acceleration[i] = 0;
            newTrajectory->endAngles[i] = start[i];
        }
        return;
    }
    float aim_velocity = max_velocity_;
    if (max_angle < max_accelerate_angle_) {
        newTrajectory->accelerate_time_step = static_cast<int>(sqrt(max_angle / max_acceleration_)) + 1;
        aim_velocity = max_angle / newTrajectory->accelerate_time_step;
        newTrajectory->uniform_time_step = 0;
    } else {
        newTrajectory->accelerate_time_step = max_accelerate_time_;
        newTrajectory->uniform_time_step = static_cast<int>((max_angle - max_accelerate_angle_) / max_velocity_);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->aim_velocity[i] = aim_velocity * angle_diff[i] / max_angle;
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / newTrajectory->accelerate_time_step;
    }
}

void TrapezoidalTrajectory::computeSteppingTrajectory(Trajectory *newTrajectory, const float *start, const float *speed)
{
    float maxSpeed = 0.0;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        if (fabs(speed[i]) > max_velocity_) {
            newTrajectory->aim_velocity[i] = (speed[i] > 0 ? 1 : -1) * max_velocity_;
        } else {
            newTrajectory->aim_velocity[i] = speed[i];
        }
        if (fabs(newTrajectory->aim_velocity[i]) > maxSpeed) {
            maxSpeed = fabs(newTrajectory->aim_velocity[i]);
        }
    }
    newTrajectory->uniform_time_step = -1;
    newTrajectory->accelerate_time_step = static_cast<int>(maxSpeed / max_acceleration_) + 1;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / newTrajectory->accelerate_time_step;
        newTrajectory->endAngles[i] = 0;
    }
}

void TrapezoidalTrajectory::computePoseSteppingTrajectory(Trajectory *newTrajectory, const float *start, const float *end)
{
    float max_angle = 0.0;
    float angle_diff[USED_ARM_DOF];
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        newTrajectory->endAngles[i] = end[i];
        angle_diff[i] = end[i] - start[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
        }
    }
    if (max_angle < 0.01) {
        newTrajectory->accelerate_time_step = 0;
        newTrajectory->uniform_time_step = 0;
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            newTrajectory->aim_velocity[i] = 0;
            newTrajectory->now_acceleration[i] = 0;
            newTrajectory->endAngles[i] = start[i];
        }
        return;
    }
    float aim_velocity = max_velocity_;
    newTrajectory->accelerate_time_step = max_accelerate_time_;
    if (max_angle < min_angle_for_pose_stepping_max_speed_ - 0.5 * max_accelerate_angle_) {
        aim_velocity = max_angle / (ARM_FOLLOWING_CHECKING_STEP * 2 - (max_accelerate_time_ - 1) * 0.5);
        newTrajectory->uniform_time_step = ARM_FOLLOWING_CHECKING_STEP * 2 - max_accelerate_time_;
    } else {
        newTrajectory->uniform_time_step = static_cast<int>((max_angle - (max_accelerate_angle_ - 1) * 0.5) / max_velocity_);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->aim_velocity[i] = aim_velocity * angle_diff[i] / max_angle;
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / max_accelerate_time_;
    }
}

int TrapezoidalTrajectory::computeTrajectory(const float *start, const float *end, int type)
{
    Trajectory* newTrajectory = addTrajectory();
    newTrajectory->goalType = type;
    newTrajectory->startTime = 0;
    newTrajectory->stopTime = 0;
    newTrajectory->current_time_step = 0;
    if (type == GOAL_TYPE_MOVING) {
        computeMovingTrajectory(newTrajectory, start, end);
    } else if (type == GOAL_TYPE_POSE_STEPPING) {
        computePoseSteppingTrajectory(newTrajectory, start, end);
    } else {
        computeSteppingTrajectory(newTrajectory, start, end);
    }
    return 0;
}

int TrapezoidalTrajectory::checkStartTrajectory(int64_t startTime)
{
    if (nowTrajectoryIndex_ < 0) {
        return -1;
    }
    if (trajectoryHistory_[nowTrajectoryIndex_].startTime != 0) {
        return 0;
    }
    trajectoryHistory_[nowTrajectoryIndex_].startTime = startTime;
    trajectoryHistory_[nowTrajectoryIndex_].stopTime = startTime
        + (2 * trajectoryHistory_[nowTrajectoryIndex_].accelerate_time_step
            + trajectoryHistory_[nowTrajectoryIndex_].uniform_time_step) * ARM_FOLLOWING_PERIOD;
    return 0;
}

int TrapezoidalTrajectory::updateTrajectory(float *aimAngles)
{
    // Update the aimAngles based on the computed trajectory
    Trajectory* newTrajectory = &trajectoryHistory_[nowTrajectoryIndex_];
    newTrajectory->current_time_step++;
    if (newTrajectory->current_time_step < newTrajectory->accelerate_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            float v = newTrajectory->now_acceleration[i] * newTrajectory->current_time_step;
            aimAngles[i] += v;
        }
    } else if (newTrajectory->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP ||
                newTrajectory->current_time_step < newTrajectory->accelerate_time_step + newTrajectory->uniform_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] += newTrajectory->aim_velocity[i];
        }
    } else if (newTrajectory->current_time_step < 2 * newTrajectory->accelerate_time_step + newTrajectory->uniform_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            int decelerate_time_step = newTrajectory->current_time_step - newTrajectory->accelerate_time_step
                - newTrajectory->uniform_time_step;
            float v = newTrajectory->aim_velocity[i] - newTrajectory->now_acceleration[i] * decelerate_time_step;
            if (v < 0) v = 0;
            aimAngles[i] += v;
        }
    } else {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] = newTrajectory->endAngles[i];
        }
        return TRAJECTORY_COMPELETE;
    }
    return TRAJECTORY_ONGOING;
}

int TrapezoidalTrajectory::getAngleFromTime(float *aimAngles, int64_t timeStamp)
{
    if (nowTrajectoryIndex_ == -1) {
        return OK;
    }
    int startIndex = nowTrajectoryIndex_ + 1;
    int endIndex = trajectoryFullFlag_ == 1 ? startIndex : 0;
    do {
        startIndex = (startIndex + MAX_TRAJECTORY_HISTORY - 1) % MAX_TRAJECTORY_HISTORY;
        Trajectory* traj = &trajectoryHistory_[startIndex];
        if (traj -> stopTime == 0) {
            if (startIndex == nowTrajectoryIndex_) {
                return ARM_NOW_FORCE_MOVING;
            }
            return UNKNOWN_ERR;
        }
        if (timeStamp >= traj->stopTime) {
            for (int i = 0; i < USED_ARM_DOF; ++i) {
                aimAngles[i] = traj->endAngles[i];
            }
            return 0;
        }
        if (timeStamp >= traj->startTime) {
            int64_t time_diff = timeStamp - traj->startTime;
            int time_step = static_cast<int>(time_diff / ARM_FOLLOWING_PERIOD);
            for (int i = 0; i < USED_ARM_DOF; ++i) {
                float angle = traj->startAngles[i];
                if (time_step < traj->accelerate_time_step) {
                    angle += 0.5 * traj->now_acceleration[i] * time_step * time_step;
                } else if (traj->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP
                        || time_step < traj->accelerate_time_step + traj->uniform_time_step) {
                    angle += 0.5 * traj->now_acceleration[i] * traj->accelerate_time_step * (traj->accelerate_time_step);
                    angle += traj->aim_velocity[i] * (time_step - traj->accelerate_time_step);
                } else if (time_step <= 2 * traj->accelerate_time_step + traj->uniform_time_step) {
                    int decelerate_time_step = time_step - traj->accelerate_time_step - traj->uniform_time_step;
                    angle += 0.5 * traj->now_acceleration[i] * (traj->accelerate_time_step) * (traj->accelerate_time_step);
                    angle += traj->aim_velocity[i] * (traj->uniform_time_step);
                    angle += traj->aim_velocity[i] * decelerate_time_step
                        - 0.5 * traj->now_acceleration[i] * decelerate_time_step * decelerate_time_step;
                } else {
                    angle = traj->endAngles[i];
                }
                aimAngles[i] = angle;
            }
            return 0;
        }
    } while (startIndex != endIndex);
    return -1;
}

static void getAngleForCheckingInner(Trajectory* traj, int timeStepOffset, float *aimAngles)
{
    int time_step = traj->current_time_step + timeStepOffset;
    if (time_step >= traj->accelerate_time_step * 2 + traj->uniform_time_step 
            && traj->uniform_time_step != STEPPING_UNIFORM_TIME_STAMP) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] = traj->endAngles[i];
        }
        return;
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        float angle = traj->startAngles[i];
        if (time_step < traj->accelerate_time_step) {
            angle += 0.5 * traj->now_acceleration[i] * time_step * time_step;
        } else if (time_step < traj->accelerate_time_step + traj->uniform_time_step || 
                    traj->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP) {
            angle += 0.5 * traj->now_acceleration[i] * (traj->accelerate_time_step) * (traj->accelerate_time_step);
            angle += traj->aim_velocity[i] * (time_step - traj->accelerate_time_step);
        } else if (time_step <= 2 * traj->accelerate_time_step + traj->uniform_time_step) {
            int decelerate_time_step = time_step - traj->accelerate_time_step - traj->uniform_time_step;
            angle += 0.5 * traj->now_acceleration[i] * (traj->accelerate_time_step) * (traj->accelerate_time_step);
            angle += traj->aim_velocity[i] * (traj->uniform_time_step);
            angle += traj->aim_velocity[i] * decelerate_time_step
                - 0.5 * traj->now_acceleration[i] * decelerate_time_step * decelerate_time_step;
        }
        aimAngles[i] = angle;
    }
}

void TrapezoidalTrajectory::getAngleForChecking(float *aimAngles)
{
    if (nowTrajectoryIndex_ < 0) {
        return;
    }
    Trajectory* traj = &trajectoryHistory_[nowTrajectoryIndex_];
    getAngleForCheckingInner(traj, ARM_FOLLOWING_CHECKING_STEP * 2, aimAngles);
}

void TrapezoidalTrajectory::getAngleForStarting(float *aimAngles)
{
    if (nowTrajectoryIndex_ < 0) {
        return;
    }
    Trajectory* traj = &trajectoryHistory_[nowTrajectoryIndex_];
    getAngleForCheckingInner(traj, ARM_FOLLOWING_CHECKING_STEP, aimAngles);
}

void TrapezoidalTrajectory::stopTrajectory()
{
    if (nowTrajectoryIndex_ < 0) {
        return;
    }
    Trajectory* traj = &trajectoryHistory_[nowTrajectoryIndex_];
    int formerPoseSteppingUniformTimeStep = traj->uniform_time_step;
    if (traj->current_time_step < traj->accelerate_time_step) {
        traj->accelerate_time_step = traj->current_time_step;
        traj->uniform_time_step = 0;
    } else if (traj->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP ||
        traj->current_time_step < traj->accelerate_time_step + traj->uniform_time_step) {
        traj->uniform_time_step = traj->current_time_step - traj->accelerate_time_step;
    }
    traj->stopTime = traj->startTime + ARM_FOLLOWING_PERIOD * (traj->accelerate_time_step * 2 + traj->uniform_time_step);
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        if (traj->goalType == GOAL_TYPE_POSE_STEPPING) {
            traj->now_acceleration[i] = traj->aim_velocity[i] / traj->accelerate_time_step;
            traj->endAngles[i] = traj->endAngles[i]
                - traj->aim_velocity[i] * (formerPoseSteppingUniformTimeStep - traj->uniform_time_step);
            traj->endAngles[i] += 0.5 * traj->now_acceleration[i] * traj->accelerate_time_step * traj->accelerate_time_step;
        } else {
            traj->endAngles[i] = traj->startAngles[i] + traj->aim_velocity[i] * (traj->uniform_time_step + traj->accelerate_time_step);
        }
    }
    if (traj->goalType == GOAL_TYPE_POSE_STEPPING) {
        traj->goalType = GOAL_TYPE_MOVING;
    }
}

int TrapezoidalTrajectory::continuePoseSteppingTrajectory(float *nowAngles, float *aimAngles)
{
    // Update the aimAngles based on the computed trajectory
    Trajectory* newTrajectory = &trajectoryHistory_[nowTrajectoryIndex_];
    float angle_diff[USED_ARM_DOF];
    float max_angle = 0.0;
    float max_velocity = 0.0;
    int fastestIndex = 0;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->endAngles[i] = aimAngles[i];
        angle_diff[i] = newTrajectory->endAngles[i] - nowAngles[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
            fastestIndex = i;
        }
    }
    max_velocity = fabs(newTrajectory->aim_velocity[fastestIndex]);
    if (max_velocity * ARM_FOLLOWING_CHECKING_STEP * 2 >= max_angle) {
        max_velocity = max_angle / (ARM_FOLLOWING_CHECKING_STEP * 2);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        if (max_angle < 0.0001) {
            newTrajectory->aim_velocity[i] = 0;
        } else {
            newTrajectory->aim_velocity[i] = max_velocity * angle_diff[i] / max_angle;
        }
    }
    newTrajectory->uniform_time_step += max_angle / max_velocity;
    return 0;
}

int TrapezoidalTrajectory::updatePoseSteppingTrajectory(float *aimAngles)
{
    Trajectory* newTrajectory = &trajectoryHistory_[nowTrajectoryIndex_];
    newTrajectory->current_time_step++;
    if (newTrajectory->current_time_step < newTrajectory->accelerate_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            float v = newTrajectory->now_acceleration[i] * newTrajectory->current_time_step;
            aimAngles[i] += v;
        }
    } else if (newTrajectory->current_time_step < newTrajectory->accelerate_time_step + newTrajectory->uniform_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] += newTrajectory->aim_velocity[i];
        }
    } else {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] += newTrajectory->aim_velocity[i];
        }
        return TRAJECTORY_COMPELETE;
    }
    return TRAJECTORY_ONGOING;
}