#include "track_core/pid_runtime.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <limits>

namespace track_core {
namespace {

constexpr float default_pid_nominal_sample_interval_s = 2.0F;

template <class... Ts>
struct overloads : Ts... {
    using Ts::operator()...;
};

bool supports_live_pid_tuning(const TrackPidConfig &config) {
    return config.schemas.size() == 1 &&
           std::holds_alternative<TrackPidSegment>(config.schemas.front().segment);
}

} // namespace

PidHrRuntime::PidHrRuntime() = default;

PidHrRuntime::PidHrRuntime(TrackPidConfig config)
    : config_(std::move(config)) {}

void PidHrRuntime::set_pid_config(TrackPidConfig config) {
    config_ = std::move(config);
}

void PidHrRuntime::update_target_hr_bpm(std::uint8_t target_hr_bpm) {
    config_.target_hr_bpm = target_hr_bpm;
    if (program_state_) {
        program_state_->update_target_hr_bpm(target_hr_bpm);
    }
}

expected<unit, TrackError> PidHrRuntime::apply_pid_runtime_command(
    const TrackPidRuntimeCommand &command,
    const TrackConfig *track_config) {
    if (!running_ || !program_state_) {
        return unexpected<TrackError>{TrackError::invalid_state};
    }

    return std::visit(overloads{
                          [](unit) -> expected<unit, TrackError> {
                              return unexpected<TrackError>{TrackError::invalid_arg};
                          },
                          [&](const TrackPidSetTargetHr &target_hr) -> expected<unit, TrackError> {
                              update_target_hr_bpm(target_hr.target_hr_bpm);
                              return {};
                          },
                          [&](const TrackPidSetTuning &tuning) -> expected<unit, TrackError> {
                              if (!supports_live_pid_tuning(config_)) {
                                  return unexpected<TrackError>{TrackError::invalid_state};
                              }
                              auto updated = program_state_->update_tuning(tuning);
                              if (!updated) {
                                  return unexpected<TrackError>{updated.error()};
                              }
                              auto &segment = std::get<TrackPidSegment>(config_.schemas.front().segment);
                              tuning.apply_to(segment);
                              base_speed_m_s_ = program_state_->commanded_speed_m_s();
                              effective_speed_m_s_ = effective_speed_m_s(base_speed_m_s_, track_config);
                              return {};
                          }},
                      command.command);
}

const TrackPidConfig &PidHrRuntime::pid_config() const {
    return config_;
}

void PidHrRuntime::start(time_point now) {
    running_ = true;
    mileage_m_ = 0.0F;
    base_speed_m_s_ = 0.0F;
    effective_speed_m_s_ = 0.0F;
    start_timestamp_ = now;
    last_tick_timestamp_ = now;
    last_consumed_hr_sample_seq_.reset();
    last_consumed_hr_sample_time_ = {};
    program_state_ = std::make_unique<TrackPidProgramState>(now, config_);
}

void PidHrRuntime::stop() {
    running_ = false;
    base_speed_m_s_ = 0.0F;
    effective_speed_m_s_ = 0.0F;
    last_consumed_hr_sample_seq_.reset();
    last_consumed_hr_sample_time_ = {};
    program_state_.reset();
}

void PidHrRuntime::tick(
    const TrackConfig *track_config,
    const TrackPidBandSnapshot &band,
    time_point now) {
    if (!running_ || !program_state_) {
        return;
    }

    const auto dt = std::chrono::duration_cast<std::chrono::duration<float>>(now - last_tick_timestamp_);
    last_tick_timestamp_ = now;
    if (dt.count() <= 0.0F) {
        return;
    }

    std::optional<TrackPidProgramState::HrSample> fresh_hr;
    if (band.hr_is_fresh && band.has_heart_rate &&
        (!last_consumed_hr_sample_seq_.has_value() ||
         band.heart_rate_sample_seq != *last_consumed_hr_sample_seq_)) {
        auto sample_interval_s = default_pid_nominal_sample_interval_s;
        if (last_consumed_hr_sample_seq_.has_value()) {
            sample_interval_s = std::chrono::duration_cast<std::chrono::duration<float>>(
                                    now - last_consumed_hr_sample_time_)
                                    .count();
        }
        fresh_hr = TrackPidProgramState::HrSample{
            .heart_rate_bpm = static_cast<float>(band.heart_rate),
            .sample_interval_s = sample_interval_s,
        };
        last_consumed_hr_sample_seq_ = band.heart_rate_sample_seq;
        last_consumed_hr_sample_time_ = now;
    }

    base_speed_m_s_ = program_state_->next(now, fresh_hr);
    effective_speed_m_s_ = effective_speed_m_s(base_speed_m_s_, track_config);
    if (program_state_->is_finished(now)) {
        base_speed_m_s_ = 0.0F;
        effective_speed_m_s_ = 0.0F;
        running_ = false;
    }

    mileage_m_ += effective_speed_m_s_ * dt.count();
}

float PidHrRuntime::effective_speed_m_s(float base_speed_m_s, const TrackConfig *track_config) const {
    auto effective_speed = std::max(0.0F, base_speed_m_s);
    if (!config_.speed_suppression || effective_speed <= 0.0F || track_config == nullptr) {
        return effective_speed;
    }

    if (!std::isfinite(track_config->line_length_m) || track_config->line_length_m <= 0.0F) {
        return effective_speed;
    }

    const auto &suppression = *config_.speed_suppression;
    if (!std::isfinite(suppression.sigma_m) || suppression.sigma_m <= 0.0F) {
        return effective_speed;
    }

    const auto phase_raw_m = std::fmod(mileage_m_, track_config->line_length_m);
    const auto phase_m = phase_raw_m < 0.0F ? (phase_raw_m + track_config->line_length_m) : phase_raw_m;
    const auto seam_distance_m = std::min(phase_m, track_config->line_length_m - phase_m);
    const auto sigma_ratio = seam_distance_m / suppression.sigma_m;
    const auto gaussian = std::exp(-0.5F * sigma_ratio * sigma_ratio);
    const auto speed_ratio = 1.0F - ((1.0F - suppression.ratio_min) * gaussian);
    return effective_speed * speed_ratio;
}

TrackReport PidHrRuntime::state_report(time_point now) const {
    return {
        .id = magic_pid_track_id,
        .state = running_ ? TrackState::run : TrackState::stop,
        .mileage_m = mileage_m_,
        .speed_m_s = effective_speed_m_s_,
        .time_elapsed_ms = static_cast<std::uint32_t>(std::chrono::duration_cast<std::chrono::milliseconds>(
                                                          now - start_timestamp_)
                                                          .count()),
    };
}

TrackInfo PidHrRuntime::info() const {
    return {
        .kind = SchemeKind::speed_input_time_segmented_mileage_free,
        .color = color_,
        .id = magic_pid_track_id,
        .is_running = running_,
        .num_segments = static_cast<std::uint8_t>(
            std::min(config_.schemas.size(), static_cast<std::size_t>(std::numeric_limits<std::uint8_t>::max()))),
    };
}

TrackPidStatus PidHrRuntime::pid_status(const TrackPidBandSnapshot &band, time_point now) const {
    TrackPidStatus status;
    status.band_id = config_.band_id;
    status.band_is_active = band.band_is_active;
    status.is_heart_rate_valid = band.hr_is_fresh;
    status.heart_rate_bpm = band.has_heart_rate ? band.heart_rate : static_cast<std::uint8_t>(0);
    status.step_count = band.has_step_count ? band.step_count : static_cast<std::uint16_t>(0);
    status.effective_speed_m_s = effective_speed_m_s_;
    status.base_speed_m_s = base_speed_m_s_;
    if (program_state_) {
        status.active_segment_index = program_state_->active_stage_index();
        status.active_segment_kind = program_state_->active_stage_kind();
        status.remaining_program_ms = program_state_->remaining_program_ms(now);
    } else {
        status.active_segment_kind = TrackPidStageKind::none;
    }
    return status;
}

} // namespace track_core