#include "bscan_processor.hpp"

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string_view>
#include <utility>
#include <vector>

namespace radar::processing {
namespace {

constexpr double kPi = 3.14159265358979323846;
constexpr double kSpeedOfLightMetersPerSec = 299'792'458.0;

struct BScanProfile {
    std::vector<float> depth_m{};
    std::vector<float> response{};
};

[[nodiscard]] auto next_power_of_two(std::size_t value) -> std::size_t {
    if (value <= 1U) {
        return 1U;
    }

    std::size_t power = 1U;
    while (power < value) {
        if (power > (std::numeric_limits<std::size_t>::max() >> 1U)) {
            return value;
        }
        power <<= 1U;
    }
    return power;
}

void fft_inplace(std::vector<std::complex<double>>& values, bool inverse) {
    const std::size_t size = values.size();
    if (size <= 1U) {
        return;
    }

    for (std::size_t index = 1U, bit_reversed = 0U; index < size; ++index) {
        std::size_t bit = size >> 1U;
        while (bit_reversed & bit) {
            bit_reversed ^= bit;
            bit >>= 1U;
        }
        bit_reversed ^= bit;
        if (index < bit_reversed) {
            std::swap(values[index], values[bit_reversed]);
        }
    }

    for (std::size_t len = 2U; len <= size; len <<= 1U) {
        const double angle = 2.0 * kPi * (inverse ? 1.0 : -1.0) / static_cast<double>(len);
        const std::complex<double> twiddle_step(std::cos(angle), std::sin(angle));
        const std::size_t half_len = len >> 1U;

        for (std::size_t offset = 0U; offset < size; offset += len) {
            std::complex<double> twiddle(1.0, 0.0);
            for (std::size_t i = 0U; i < half_len; ++i) {
                const auto even = values[offset + i];
                const auto odd = values[offset + i + half_len] * twiddle;

                values[offset + i] = even + odd;
                values[offset + i + half_len] = even - odd;
                twiddle *= twiddle_step;
            }
        }
    }

    if (!inverse) {
        return;
    }

    const double scale = 1.0 / static_cast<double>(size);
    for (auto& value : values) {
        value *= scale;
    }
}

[[nodiscard]] auto select_axis(std::complex<double> sample, std::string_view axis) -> double {
    if (axis == "real") {
        return std::real(sample);
    }
    if (axis == "phase") {
        return std::arg(sample);
    }
    return std::abs(sample);
}

[[nodiscard]] auto fallback_profile(const ipc::SweepTraceBlock& trace) -> BScanProfile {
    const std::size_t point_count = std::min(trace.frequency_hz.size(), trace.s21.size());
    BScanProfile fallback{};
    fallback.depth_m.reserve(point_count);
    fallback.response.reserve(point_count);

    const float denominator = point_count > 1U ? static_cast<float>(point_count - 1U) : 1.0F;
    for (std::size_t index = 0U; index < point_count; ++index) {
        const auto& sample = trace.s21[index];
        fallback.depth_m.push_back(static_cast<float>(static_cast<float>(index) / denominator));
        fallback.response.push_back(std::sqrt(sample.re * sample.re + sample.im * sample.im));
    }
    return fallback;
}

[[nodiscard]] auto compute_bscan_profile(
    const ipc::SweepTraceBlock& trace,
    const ProcessingLiveConfig& live_config
) -> BScanProfile {
    const std::size_t point_count = std::min(trace.frequency_hz.size(), trace.s21.size());
    if (point_count < 2U) {
        return fallback_profile(trace);
    }

    const double configured_start_hz = static_cast<double>(live_config.bscan_start_freq_mhz) * 1'000'000.0;
    const double configured_stop_hz = static_cast<double>(live_config.bscan_stop_freq_mhz) * 1'000'000.0;
    const double start_hz = std::min(configured_start_hz, configured_stop_hz);
    const double stop_hz = std::max(configured_start_hz, configured_stop_hz);

    std::vector<double> filtered_freq_hz{};
    std::vector<std::complex<double>> filtered_s21{};
    filtered_freq_hz.reserve(point_count);
    filtered_s21.reserve(point_count);

    for (std::size_t index = 0U; index < point_count; ++index) {
        const double frequency_hz = static_cast<double>(trace.frequency_hz[index]);
        if (frequency_hz < start_hz || frequency_hz > stop_hz) {
            continue;
        }
        const auto& sample = trace.s21[index];
        filtered_freq_hz.push_back(frequency_hz);
        filtered_s21.emplace_back(static_cast<double>(sample.re), static_cast<double>(sample.im));
    }

    if (filtered_freq_hz.size() < 2U) {
        return fallback_profile(trace);
    }

    const std::size_t filtered_count = filtered_freq_hz.size();
    const double df = (filtered_freq_hz.back() - filtered_freq_hz.front()) / static_cast<double>(filtered_count - 1U);
    if (df <= 0.0) {
        return fallback_profile(trace);
    }

    const auto start_bin = static_cast<std::int64_t>(std::llround(filtered_freq_hz.front() / df));
    if (start_bin < 0) {
        return fallback_profile(trace);
    }

    const auto start_index = static_cast<std::size_t>(start_bin);
    if (start_index > (std::numeric_limits<std::size_t>::max() / 2U)) {
        return fallback_profile(trace);
    }
    if (start_index > (std::numeric_limits<std::size_t>::max() - filtered_count + 1U)) {
        return fallback_profile(trace);
    }

    const std::size_t min_fft_len = 2U * (start_index + filtered_count - 1U);
    const std::size_t fft_len = next_power_of_two(min_fft_len);
    if (fft_len < min_fft_len || (fft_len & (fft_len - 1U)) != 0U) {
        return fallback_profile(trace);
    }
    if (start_index > fft_len || filtered_count > (fft_len - start_index)) {
        return fallback_profile(trace);
    }

    std::vector<std::complex<double>> spectrum(fft_len, std::complex<double>(0.0, 0.0));
    for (std::size_t index = 0U; index < filtered_count; ++index) {
        spectrum[start_index + index] = filtered_s21[index];
    }

    fft_inplace(spectrum, true);

    const double dt = 1.0 / (static_cast<double>(fft_len) * df);
    const double cut_m = std::max(0.0, static_cast<double>(live_config.bscan_cut_m));
    const double max_depth_m = std::max(0.0, static_cast<double>(live_config.bscan_max_depth_m));
    const double gain = static_cast<double>(live_config.bscan_gain);
    const double window_start = 2.0 * cut_m;
    const double window_stop = window_start + (2.0 * max_depth_m);

    BScanProfile profile{};
    profile.depth_m.reserve(spectrum.size());
    profile.response.reserve(spectrum.size());

    const std::string_view axis = live_config.bscan_axis;
    for (std::size_t index = 0U; index < spectrum.size(); ++index) {
        const double depth_raw = static_cast<double>(index) * dt * kSpeedOfLightMetersPerSec;
        if (depth_raw < window_start || depth_raw > window_stop) {
            continue;
        }

        const double one_way_depth = (depth_raw - window_start) / 2.0;
        double gain_shape = std::pow(one_way_depth, gain);
        if (!std::isfinite(gain_shape)) {
            gain_shape = 0.0;
        }

        const double axis_value = select_axis(spectrum[index], axis);
        const double output_value = axis_value * gain_shape;

        profile.depth_m.push_back(static_cast<float>(one_way_depth));
        profile.response.push_back(static_cast<float>(output_value));
    }

    return profile;
}

}  // namespace

auto BScanProcessor::name() const -> std::string {
    return "bscan";
}

auto BScanProcessor::process(const ipc::SweepTraceBlock& trace, const ProcessingLiveConfig& live_config)
    -> ipc::ResultPayload {
    ipc::ResultPayload payload{};
    payload.processing_name = name();
    payload.kind = ipc::ResultKind::TraceComplex;

    auto profile = compute_bscan_profile(trace, live_config);
    payload.frequency_hz = std::move(profile.depth_m);
    payload.trace.reserve(profile.response.size());
    for (const auto value : profile.response) {
        payload.trace.push_back(ipc::Complex32{
            .re = value,
            .im = 0.0F,
        });
    }

    return payload;
}

}  // namespace radar::processing
