#include "rm_control/modules/trajectory/drag_trajectory_validator.h"

#include "rm_control/common/rm75_constants.h"

#include <cmath>
#include <cstddef>
#include <sstream>
#include <string>
#include <vector>

namespace rm_control::modules::trajectory {
namespace {

constexpr double kMillisecondsPerSecond = 1000.0;

common::Result<void> ValidateLimits(const DragTrajectoryValidationLimits &limits) {
    const bool jump_is_valid = std::isfinite(limits.maximum_joint_jump_deg) &&
                               limits.maximum_joint_jump_deg > 0.0;
    const bool velocity_is_valid =
        std::isfinite(limits.maximum_joint_velocity_deg_per_second) &&
        limits.maximum_joint_velocity_deg_per_second > 0.0;
    const bool acceleration_is_valid =
        std::isfinite(limits.maximum_joint_acceleration_deg_per_second_squared) &&
        limits.maximum_joint_acceleration_deg_per_second_squared > 0.0;

    if (!jump_is_valid || !velocity_is_valid || !acceleration_is_valid) {
        return common::Result<void>::failure(
            "trajectory validation limits must be finite positive values");
    }
    return common::Result<void>::success();
}

bool PoseIsFinite(const common::Pose &pose) {
    return std::isfinite(pose.x) && std::isfinite(pose.y) && std::isfinite(pose.z) &&
           std::isfinite(pose.rx) && std::isfinite(pose.ry) && std::isfinite(pose.rz);
}

common::Result<void> ValidateSampleShapeAndValues(const common::TrajectorySample &sample,
                                                  std::size_t sample_index) {
    if (sample.sample_index != sample_index) {
        return common::Result<void>::failure(
            "trajectory sample_index must be contiguous and start at zero");
    }
    if (sample.source.empty()) {
        return common::Result<void>::failure(
            "trajectory sample " + std::to_string(sample_index) +
            " must contain a non-empty source");
    }
    if (sample.joints_deg.size() != common::kRm75JointCount) {
        return common::Result<void>::failure(
            "trajectory sample " + std::to_string(sample_index) +
            " must contain exactly 7 RM75 joints");
    }
    for (std::size_t joint_index = 0; joint_index < common::kRm75JointCount; ++joint_index) {
        if (!std::isfinite(sample.joints_deg[joint_index])) {
            return common::Result<void>::failure(
                "trajectory sample " + std::to_string(sample_index) + " joint " +
                std::to_string(joint_index) + " is not finite");
        }
    }
    if (!PoseIsFinite(sample.tcp_pose)) {
        return common::Result<void>::failure(
            "trajectory sample " + std::to_string(sample_index) + " TCP pose is not finite");
    }
    return common::Result<void>::success();
}

common::Result<void> ValidateSamplesBeforeKinematics(
    const std::vector<common::TrajectorySample> &samples) {
    if (samples.size() < 2U) {
        return common::Result<void>::failure(
            "trajectory must contain at least two samples");
    }
    if (samples.front().elapsed_ms < 0) {
        return common::Result<void>::failure(
            "trajectory first elapsed timestamp must not be negative");
    }

    for (std::size_t sample_index = 0; sample_index < samples.size(); ++sample_index) {
        const auto value_result = ValidateSampleShapeAndValues(samples[sample_index], sample_index);
        if (!value_result.ok) {
            return value_result;
        }
        if (sample_index > 0 &&
            samples[sample_index].elapsed_ms <= samples[sample_index - 1].elapsed_ms) {
            return common::Result<void>::failure(
                "trajectory timestamps must be strictly increasing at sample " +
                std::to_string(sample_index));
        }
    }
    return common::Result<void>::success();
}

std::string LimitExceededMessage(const std::string &metric_name,
                                 std::size_t sample_index,
                                 std::size_t joint_index,
                                 double observed_value,
                                 double allowed_value) {
    std::ostringstream message;
    message << metric_name << " limit exceeded at sample " << sample_index << ", joint "
            << joint_index << ": observed " << observed_value << ", allowed " << allowed_value;
    return message.str();
}

common::Result<std::vector<double>> MeasureIntervalVelocities(
    const common::TrajectorySample &previous,
    const common::TrajectorySample &current,
    std::size_t current_sample_index,
    const DragTrajectoryValidationLimits &limits,
    DragTrajectoryValidationReport *report) {
    const double interval_seconds =
        static_cast<double>(current.elapsed_ms - previous.elapsed_ms) / kMillisecondsPerSecond;
    std::vector<double> velocities(common::kRm75JointCount, 0.0);

    for (std::size_t joint_index = 0; joint_index < common::kRm75JointCount; ++joint_index) {
        const double joint_jump_deg = std::abs(
            static_cast<double>(current.joints_deg[joint_index]) -
            static_cast<double>(previous.joints_deg[joint_index]));
        const double velocity_deg_per_second = joint_jump_deg / interval_seconds;
        velocities[joint_index] =
            (static_cast<double>(current.joints_deg[joint_index]) -
             static_cast<double>(previous.joints_deg[joint_index])) /
            interval_seconds;

        if (joint_jump_deg > report->observed_maximum_joint_jump_deg) {
            report->observed_maximum_joint_jump_deg = joint_jump_deg;
            report->maximum_jump_sample_index = current_sample_index;
            report->maximum_jump_joint_index = joint_index;
        }
        if (velocity_deg_per_second > report->observed_maximum_joint_velocity_deg_per_second) {
            report->observed_maximum_joint_velocity_deg_per_second = velocity_deg_per_second;
            report->maximum_velocity_sample_index = current_sample_index;
            report->maximum_velocity_joint_index = joint_index;
        }

        if (joint_jump_deg > limits.maximum_joint_jump_deg) {
            return common::Result<std::vector<double>>::failure(LimitExceededMessage(
                "joint jump", current_sample_index, joint_index, joint_jump_deg,
                limits.maximum_joint_jump_deg));
        }
        if (velocity_deg_per_second > limits.maximum_joint_velocity_deg_per_second) {
            return common::Result<std::vector<double>>::failure(LimitExceededMessage(
                "joint velocity", current_sample_index, joint_index, velocity_deg_per_second,
                limits.maximum_joint_velocity_deg_per_second));
        }
    }
    return common::Result<std::vector<double>>::success(velocities);
}

common::Result<void> MeasureIntervalAcceleration(
    const std::vector<double> &previous_velocities,
    const std::vector<double> &current_velocities,
    double previous_interval_seconds,
    double current_interval_seconds,
    std::size_t current_sample_index,
    const DragTrajectoryValidationLimits &limits,
    DragTrajectoryValidationReport *report) {
    const double interval_midpoint_delta_seconds =
        (previous_interval_seconds + current_interval_seconds) / 2.0;
    for (std::size_t joint_index = 0; joint_index < common::kRm75JointCount; ++joint_index) {
        const double acceleration_deg_per_second_squared =
            std::abs(current_velocities[joint_index] - previous_velocities[joint_index]) /
            interval_midpoint_delta_seconds;
        if (acceleration_deg_per_second_squared >
            report->observed_maximum_joint_acceleration_deg_per_second_squared) {
            report->observed_maximum_joint_acceleration_deg_per_second_squared =
                acceleration_deg_per_second_squared;
            report->maximum_acceleration_sample_index = current_sample_index;
            report->maximum_acceleration_joint_index = joint_index;
        }
        if (acceleration_deg_per_second_squared >
            limits.maximum_joint_acceleration_deg_per_second_squared) {
            return common::Result<void>::failure(LimitExceededMessage(
                "joint acceleration", current_sample_index, joint_index,
                acceleration_deg_per_second_squared,
                limits.maximum_joint_acceleration_deg_per_second_squared));
        }
    }
    return common::Result<void>::success();
}

} // namespace

common::Result<DragTrajectoryValidationReport> ValidateDragTrajectory(
    const std::vector<common::TrajectorySample> &samples,
    const DragTrajectoryValidationLimits &limits) {
    const auto limits_result = ValidateLimits(limits);
    if (!limits_result.ok) {
        return common::Result<DragTrajectoryValidationReport>::failure(limits_result.error_message);
    }
    const auto samples_result = ValidateSamplesBeforeKinematics(samples);
    if (!samples_result.ok) {
        return common::Result<DragTrajectoryValidationReport>::failure(samples_result.error_message);
    }

    DragTrajectoryValidationReport report;
    report.sample_count = samples.size();
    report.duration_ms = samples.back().elapsed_ms - samples.front().elapsed_ms;
    std::vector<double> previous_velocities;
    double previous_interval_seconds = 0.0;

    for (std::size_t sample_index = 1; sample_index < samples.size(); ++sample_index) {
        const double current_interval_seconds =
            static_cast<double>(samples[sample_index].elapsed_ms -
                                samples[sample_index - 1].elapsed_ms) /
            kMillisecondsPerSecond;
        const auto velocities_result = MeasureIntervalVelocities(
            samples[sample_index - 1], samples[sample_index], sample_index, limits, &report);
        if (!velocities_result.ok) {
            return common::Result<DragTrajectoryValidationReport>::failure(
                velocities_result.error_message);
        }

        if (!previous_velocities.empty()) {
            const auto acceleration_result = MeasureIntervalAcceleration(
                previous_velocities, velocities_result.value, previous_interval_seconds,
                current_interval_seconds, sample_index, limits, &report);
            if (!acceleration_result.ok) {
                return common::Result<DragTrajectoryValidationReport>::failure(
                    acceleration_result.error_message);
            }
        }
        previous_velocities = velocities_result.value;
        previous_interval_seconds = current_interval_seconds;
    }

    return common::Result<DragTrajectoryValidationReport>::success(report);
}

} // namespace rm_control::modules::trajectory
