#include "data_preprocessor.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <span>
#include <stdexcept>
#include <string_view>
#include <thread>
#include <vector>

namespace radar::preprocessing {
namespace {

constexpr double kPi = 3.14159265358979323846;

[[nodiscard]] auto should_apply_notch(const config::PreprocessNotchConfig& notch) -> bool {
    return notch.enabled && !notch.bands_hz.empty();
}

[[nodiscard]] auto build_notch_mask(
    std::span<const float> frequency_hz,
    const config::PreprocessNotchConfig& notch
) -> std::vector<float> {
    std::vector<float> mask(frequency_hz.size(), 1.0F);
    if (!should_apply_notch(notch)) {
        return mask;
    }

    const bool hard_taper = notch.taper_type == "hard";
    const double taper_width_hz = std::max(0.0, static_cast<double>(notch.taper_width_hz));
    for (const auto& band : notch.bands_hz) {
        const double low_hz = static_cast<double>(band.low_hz);
        const double high_hz = static_cast<double>(band.high_hz);
        const double trans_low_hz = low_hz - taper_width_hz;
        const double trans_high_hz = high_hz + taper_width_hz;

        for (std::size_t index = 0U; index < frequency_hz.size(); ++index) {
            const double frequency = static_cast<double>(frequency_hz[index]);
            if (frequency >= low_hz && frequency <= high_hz) {
                mask[index] = 0.0F;
                continue;
            }

            if (hard_taper || taper_width_hz <= 0.0) {
                if (frequency >= trans_low_hz && frequency <= trans_high_hz) {
                    mask[index] = 0.0F;
                }
                continue;
            }

            if (frequency >= trans_low_hz && frequency < low_hz) {
                const double arg = kPi * (frequency - trans_low_hz) / taper_width_hz;
                mask[index] *= static_cast<float>(0.5 * (1.0 - std::cos(arg)));
                continue;
            }

            if (frequency > high_hz && frequency <= trans_high_hz) {
                const double arg = kPi * (frequency - high_hz) / taper_width_hz;
                mask[index] *= static_cast<float>(0.5 * (1.0 + std::cos(arg)));
            }
        }
    }

    return mask;
}

void apply_notch_mask(std::vector<ipc::Complex32>& samples, std::span<const float> mask) {
    const std::size_t point_count = std::min(samples.size(), mask.size());
    for (std::size_t index = 0U; index < point_count; ++index) {
        samples[index].re *= mask[index];
        samples[index].im *= mask[index];
    }
}

[[nodiscard]] auto apply_notch_filter(
    const ipc::SweepTraceBlock& trace,
    const config::PreprocessNotchConfig& notch
) -> ipc::SweepTraceBlock {
    if (!should_apply_notch(notch)) {
        return trace;
    }

    ipc::SweepTraceBlock filtered = trace;
    const auto mask = build_notch_mask(filtered.frequency_hz, notch);
    apply_notch_mask(filtered.s21, mask);
    apply_notch_mask(filtered.s11, mask);
    return filtered;
}

}  // namespace

DataPreprocessor::DataPreprocessor(
    const config::RunConfig& config,
    CalibrationMaster& calibration_master,
    ReferenceMaster& reference_master,
    ipc::ShmRing& raw_ring,
    ipc::ShmRing& preprocessed_ring,
    ipc::ShmRing* preprocessed_tap_ring
)
    : config_(config),
      calibration_master_(calibration_master),
      reference_master_(reference_master),
      raw_ring_(raw_ring),
      preprocessed_ring_(preprocessed_ring),
      preprocessed_tap_ring_(preprocessed_tap_ring) {}

void DataPreprocessor::validate_startup_prerequisites() const {
    calibration_master_.validate_combos(config_.run_combos);
    reference_master_.validate_combos(config_.run_combos);
}

auto DataPreprocessor::try_pop_raw_collection(
    std::vector<std::uint8_t>& serialized_raw,
    ipc::RawSweepCollection* out_collection
) -> bool {
    if (!raw_ring_.pop(serialized_raw)) {
        std::this_thread::sleep_for(std::chrono::milliseconds(config_.runtime.idle_sleep_ms));
        return false;
    }

    *out_collection = ipc::deserialize_raw_collection(serialized_raw);
    return true;
}

void DataPreprocessor::publish_preprocessed_collection(const ipc::PreprocessedCollection& collection) {
    const auto serialized_preprocessed = ipc::serialize_preprocessed_collection(collection);
    if (!preprocessed_ring_.push(serialized_preprocessed)) {
        throw std::runtime_error("Preprocessed ring slot is too small for serialized collection");
    }
    if (preprocessed_tap_ring_ != nullptr && !preprocessed_tap_ring_->push(serialized_preprocessed)) {
        throw std::runtime_error("Preprocessed tap ring slot is too small for serialized collection");
    }
}

void DataPreprocessor::run(const std::atomic<bool>& stop_requested) {
    validate_startup_prerequisites();

    std::vector<std::uint8_t> serialized_raw{};
    ipc::RawSweepCollection raw_collection{};
    while (!stop_requested.load(std::memory_order_relaxed)) {
        if (!try_pop_raw_collection(serialized_raw, &raw_collection)) {
            continue;
        }

        const auto preprocessed_collection = preprocess_collection(raw_collection);
        publish_preprocessed_collection(preprocessed_collection);
    }
}

auto DataPreprocessor::preprocess_collection(const ipc::RawSweepCollection& raw_collection) const
    -> ipc::PreprocessedCollection {
    ipc::PreprocessedCollection preprocessed{};
    preprocessed.collection_id = raw_collection.collection_id;
    preprocessed.monotonic_ns = ipc::current_monotonic_ns();
    preprocessed.capture_start_ns = raw_collection.capture_start_ns;
    preprocessed.capture_end_ns = raw_collection.capture_end_ns;
    preprocessed.traces.reserve(raw_collection.traces.size());

    for (const auto& raw_trace : raw_collection.traces) {
        // Pipeline order is fixed: calibration first, then reference subtraction.
        const auto calibrated = calibration_master_.apply_to_trace(raw_trace);
        const auto referenced = reference_master_.apply_to_trace(calibrated);
        preprocessed.traces.push_back(apply_notch_filter(referenced, config_.preprocess.notch));
    }

    return preprocessed;
}

}  // namespace radar::preprocessing
