#include "gpr_processor.hpp"

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <numeric>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace radar::processing {
namespace {

constexpr double kPi = 3.14159265358979323846;
constexpr double kSpeedOfLightMetersPerSec = 299'792'458.0;
constexpr double kAccumulatorXMarginM = 2.0;
constexpr double kAccumulatorZMinM = 0.10;
constexpr double kSnrThresh = 3.0;
constexpr double kSnrCompMax = 20.0;
constexpr std::size_t kGridWidth = 300U;
constexpr std::size_t kGridHeight = 300U;
constexpr double kGaussianSigma = 3.0;
constexpr double kCleanSuppressRadiusM = 0.07;
constexpr double kCleanThresholdFrac = 0.05;
constexpr std::size_t kMaxObjects = 15U;
constexpr double kExtendedThresholdFrac = 0.75;
constexpr double kExtendedMinAreaCm2 = 2.0;

using PairKey = std::uint64_t;

struct SelectedTrace {
    std::vector<double> frequency_hz{};
    std::vector<std::complex<double>> s21{};
};

struct AscanResult {
    std::vector<double> time_s{};
    std::vector<double> depth_m{};
    std::vector<double> amplitude{};
    double bandwidth_hz = 0.0;
};

struct PeakRecord {
    double z_app = 0.0;
    double tau = 0.0;
    double snr_raw = 0.0;
    double snr_comp = 0.0;
};

struct PointRecord {
    double x_m = 0.0;
    double z_m = 0.0;
    double score = 0.0;
};

struct RegionRecord {
    double x_m = 0.0;
    double z_m = 0.0;
    double score = 0.0;
    double pixel_count = 0.0;
    std::vector<float> mask{};
};

struct GeometrySelection {
    std::vector<std::uint32_t> input_positions{};
    std::vector<std::uint32_t> output_positions{};
    std::vector<double> x_tx{};
    std::vector<double> x_rx{};
    std::unordered_map<std::uint32_t, std::uint32_t> input_local_by_pos{};
    std::unordered_map<std::uint32_t, std::uint32_t> output_local_by_pos{};
};

struct BackgroundAccumulator {
    std::vector<std::complex<double>> sum{};
    std::size_t count = 0U;
};

struct GridDefinition {
    std::vector<double> x_grid{};
    std::vector<double> z_grid{};
    std::vector<std::vector<double>> tx_distance_grids{};
    std::vector<std::vector<double>> rx_distance_grids{};
};

[[nodiscard]] auto make_pair_key(std::uint32_t tx_local_index, std::uint32_t rx_local_index) -> PairKey {
    return (static_cast<PairKey>(tx_local_index) << 32U) | static_cast<PairKey>(rx_local_index);
}

[[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 inner = 0U; inner < half_len; ++inner) {
                const auto even = values[offset + inner];
                const auto odd = values[offset + inner + half_len] * twiddle;
                values[offset + inner] = even + odd;
                values[offset + inner + 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 build_axis(double min_value, double max_value, std::size_t count) -> std::vector<double> {
    std::vector<double> axis{};
    if (count == 0U) {
        return axis;
    }

    axis.resize(count, min_value);
    if (count == 1U) {
        return axis;
    }

    const double step = (max_value - min_value) / static_cast<double>(count - 1U);
    for (std::size_t index = 0U; index < count; ++index) {
        axis[index] = min_value + (step * static_cast<double>(index));
    }
    return axis;
}

[[nodiscard]] auto clamp_index(std::ptrdiff_t value, std::size_t limit) -> std::size_t {
    if (limit == 0U) {
        return 0U;
    }
    if (value < 0) {
        return 0U;
    }
    const auto max_index = static_cast<std::ptrdiff_t>(limit - 1U);
    if (value > max_index) {
        return limit - 1U;
    }
    return static_cast<std::size_t>(value);
}

[[nodiscard]] auto build_gaussian_kernel(double sigma) -> std::vector<double> {
    if (!(sigma > 0.0)) {
        return {1.0};
    }

    const auto radius = static_cast<std::ptrdiff_t>(std::ceil(sigma * 3.0));
    std::vector<double> kernel(static_cast<std::size_t>((radius * 2) + 1), 0.0);
    double sum = 0.0;
    for (std::ptrdiff_t offset = -radius; offset <= radius; ++offset) {
        const double value = std::exp(-0.5 * std::pow(static_cast<double>(offset) / sigma, 2.0));
        kernel[static_cast<std::size_t>(offset + radius)] = value;
        sum += value;
    }
    if (sum > 0.0) {
        for (auto& value : kernel) {
            value /= sum;
        }
    }
    return kernel;
}

[[nodiscard]] auto gaussian_filter_2d(
    const std::vector<double>& values,
    std::size_t width,
    std::size_t height,
    double sigma
) -> std::vector<double> {
    if (values.empty() || width == 0U || height == 0U) {
        return {};
    }

    const auto kernel = build_gaussian_kernel(sigma);
    const auto radius = static_cast<std::ptrdiff_t>((kernel.size() - 1U) / 2U);
    std::vector<double> temp(values.size(), 0.0);
    std::vector<double> output(values.size(), 0.0);

    for (std::size_t row = 0U; row < height; ++row) {
        for (std::size_t col = 0U; col < width; ++col) {
            double sum = 0.0;
            for (std::ptrdiff_t offset = -radius; offset <= radius; ++offset) {
                const auto sample_col = clamp_index(static_cast<std::ptrdiff_t>(col) + offset, width);
                sum += values[(row * width) + sample_col] * kernel[static_cast<std::size_t>(offset + radius)];
            }
            temp[(row * width) + col] = sum;
        }
    }

    for (std::size_t row = 0U; row < height; ++row) {
        for (std::size_t col = 0U; col < width; ++col) {
            double sum = 0.0;
            for (std::ptrdiff_t offset = -radius; offset <= radius; ++offset) {
                const auto sample_row = clamp_index(static_cast<std::ptrdiff_t>(row) + offset, height);
                sum += temp[(sample_row * width) + col] * kernel[static_cast<std::size_t>(offset + radius)];
            }
            output[(row * width) + col] = sum;
        }
    }

    return output;
}

[[nodiscard]] auto median_copy(std::vector<double> values) -> double {
    if (values.empty()) {
        return 0.0;
    }

    const auto middle = values.begin() + static_cast<std::ptrdiff_t>(values.size() / 2U);
    std::nth_element(values.begin(), middle, values.end());
    double median = *middle;
    if ((values.size() % 2U) == 0U) {
        const auto lower_middle = values.begin() + static_cast<std::ptrdiff_t>((values.size() / 2U) - 1U);
        std::nth_element(values.begin(), lower_middle, values.end());
        median = 0.5 * (median + *lower_middle);
    }
    return median;
}

[[nodiscard]] auto find_peak_indices(
    const std::vector<double>& values,
    std::size_t start_index,
    std::size_t stop_index,
    double threshold,
    std::size_t min_distance
) -> std::vector<std::size_t> {
    if (stop_index <= start_index + 2U) {
        return {};
    }

    std::vector<std::size_t> candidates{};
    for (std::size_t index = start_index + 1U; index + 1U < stop_index; ++index) {
        if (values[index] < threshold) {
            continue;
        }
        if (values[index] <= values[index - 1U]) {
            continue;
        }
        if (values[index] < values[index + 1U]) {
            continue;
        }
        candidates.push_back(index);
    }

    std::sort(candidates.begin(), candidates.end(), [&](const std::size_t left, const std::size_t right) {
        return values[left] > values[right];
    });

    std::vector<std::size_t> selected{};
    for (const auto index : candidates) {
        bool keep = true;
        for (const auto accepted : selected) {
            const auto distance = accepted > index ? accepted - index : index - accepted;
            if (distance < min_distance) {
                keep = false;
                break;
            }
        }
        if (keep) {
            selected.push_back(index);
        }
    }

    std::sort(selected.begin(), selected.end());
    return selected;
}

[[nodiscard]] auto attenuation_at_depth(
    std::size_t tx_index,
    std::size_t rx_index,
    double z_app,
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx
) -> double {
    const double xc = (x_tx[tx_index] + x_rx[rx_index]) / 2.0;
    const double r_tx = std::sqrt(std::pow(xc - x_tx[tx_index], 2.0) + std::pow(z_app, 2.0));
    const double r_rx = std::sqrt(std::pow(xc - x_rx[rx_index], 2.0) + std::pow(z_app, 2.0));
    const double geo = 1.0 / ((r_tx * r_rx) + 1e-12);
    const double pat =
        std::pow(z_app / (r_tx + 1e-12), 2.0) * std::pow(z_app / (r_rx + 1e-12), 2.0);
    return geo * pat + 1e-30;
}

[[nodiscard]] auto lower_bound_index(const std::vector<double>& axis, double value) -> std::size_t {
    return static_cast<std::size_t>(std::distance(axis.begin(), std::lower_bound(axis.begin(), axis.end(), value)));
}

[[nodiscard]] auto build_geometry_selection(
    const config::RunConfig& run_config,
    const ProcessingLiveConfig& live_config
) -> GeometrySelection {
    GeometrySelection selection{};

    std::unordered_map<std::uint32_t, double> tx_x_by_pos{};
    for (const auto& entry : run_config.gpr.tx_geometry) {
        tx_x_by_pos[entry.output_pos] = static_cast<double>(entry.x_m);
    }

    std::unordered_map<std::uint32_t, double> rx_x_by_pos{};
    for (const auto& entry : run_config.gpr.rx_geometry) {
        rx_x_by_pos[entry.input_pos] = static_cast<double>(entry.x_m);
    }

    auto select_positions = [](const std::vector<std::uint32_t>& requested, std::vector<std::uint32_t> available) {
        std::sort(available.begin(), available.end());
        available.erase(std::unique(available.begin(), available.end()), available.end());
        if (requested.empty()) {
            return available;
        }

        std::vector<std::uint32_t> result{};
        for (const auto value : requested) {
            if (std::find(available.begin(), available.end(), value) != available.end()) {
                result.push_back(value);
            }
        }
        std::sort(result.begin(), result.end());
        result.erase(std::unique(result.begin(), result.end()), result.end());
        return result;
    };

    std::vector<std::uint32_t> available_outputs{};
    available_outputs.reserve(tx_x_by_pos.size());
    for (const auto& [position, _] : tx_x_by_pos) {
        available_outputs.push_back(position);
    }

    std::vector<std::uint32_t> available_inputs{};
    available_inputs.reserve(rx_x_by_pos.size());
    for (const auto& [position, _] : rx_x_by_pos) {
        available_inputs.push_back(position);
    }

    selection.output_positions = select_positions(live_config.gpr_output_positions, std::move(available_outputs));
    selection.input_positions = select_positions(live_config.gpr_input_positions, std::move(available_inputs));

    selection.x_tx.reserve(selection.output_positions.size());
    for (std::size_t index = 0U; index < selection.output_positions.size(); ++index) {
        const auto position = selection.output_positions[index];
        selection.output_local_by_pos[position] = static_cast<std::uint32_t>(index);
        selection.x_tx.push_back(tx_x_by_pos[position]);
    }

    selection.x_rx.reserve(selection.input_positions.size());
    for (std::size_t index = 0U; index < selection.input_positions.size(); ++index) {
        const auto position = selection.input_positions[index];
        selection.input_local_by_pos[position] = static_cast<std::uint32_t>(index);
        selection.x_rx.push_back(rx_x_by_pos[position]);
    }

    return selection;
}

[[nodiscard]] auto build_background_mean(
    std::span<const ipc::PreprocessedCollection> previous_collections,
    const GeometrySelection& selection,
    const ProcessingLiveConfig& live_config
) -> std::unordered_map<PairKey, std::vector<std::complex<double>>> {
    std::unordered_map<PairKey, std::vector<std::complex<double>>> result{};
    if (!live_config.gpr_background_subtract_enabled || live_config.gpr_background_mean_count == 0U) {
        return result;
    }

    const std::size_t mean_count = static_cast<std::size_t>(live_config.gpr_background_mean_count);
    const std::size_t start_index =
        previous_collections.size() > mean_count ? previous_collections.size() - mean_count : 0U;

    std::unordered_map<PairKey, BackgroundAccumulator> accumulators{};
    for (std::size_t collection_index = start_index; collection_index < previous_collections.size(); ++collection_index) {
        const auto& collection = previous_collections[collection_index];
        for (const auto& trace : collection.traces) {
            const auto output_it = selection.output_local_by_pos.find(trace.combo.output_pos);
            const auto input_it = selection.input_local_by_pos.find(trace.combo.input_pos);
            if (output_it == selection.output_local_by_pos.end() || input_it == selection.input_local_by_pos.end()) {
                continue;
            }

            const auto key = make_pair_key(output_it->second, input_it->second);
            auto& accumulator = accumulators[key];
            if (accumulator.sum.empty()) {
                accumulator.sum.assign(trace.s21.size(), std::complex<double>(0.0, 0.0));
            }
            if (accumulator.sum.size() != trace.s21.size()) {
                continue;
            }

            for (std::size_t sample_index = 0U; sample_index < trace.s21.size(); ++sample_index) {
                const auto& sample = trace.s21[sample_index];
                accumulator.sum[sample_index] += std::complex<double>(sample.re, sample.im);
            }
            accumulator.count += 1U;
        }
    }

    for (auto& [key, accumulator] : accumulators) {
        if (accumulator.count == 0U) {
            continue;
        }

        auto& mean_trace = result[key];
        mean_trace = std::move(accumulator.sum);
        const double inverse_count = 1.0 / static_cast<double>(accumulator.count);
        for (auto& sample : mean_trace) {
            sample *= inverse_count;
        }
    }

    return result;
}

[[nodiscard]] auto collect_selected_traces(
    const ipc::PreprocessedCollection& collection,
    const GeometrySelection& selection,
    const std::unordered_map<PairKey, std::vector<std::complex<double>>>& background_mean
) -> std::unordered_map<PairKey, SelectedTrace> {
    std::unordered_map<PairKey, SelectedTrace> traces{};

    for (const auto& trace : collection.traces) {
        const auto output_it = selection.output_local_by_pos.find(trace.combo.output_pos);
        const auto input_it = selection.input_local_by_pos.find(trace.combo.input_pos);
        if (output_it == selection.output_local_by_pos.end() || input_it == selection.input_local_by_pos.end()) {
            continue;
        }

        SelectedTrace selected{};
        selected.frequency_hz.reserve(trace.frequency_hz.size());
        for (const auto value : trace.frequency_hz) {
            selected.frequency_hz.push_back(static_cast<double>(value));
        }

        selected.s21.reserve(trace.s21.size());
        const auto key = make_pair_key(output_it->second, input_it->second);
        const auto background_it = background_mean.find(key);
        for (std::size_t sample_index = 0U; sample_index < trace.s21.size(); ++sample_index) {
            std::complex<double> sample(trace.s21[sample_index].re, trace.s21[sample_index].im);
            if (background_it != background_mean.end() && background_it->second.size() == trace.s21.size()) {
                sample -= background_it->second[sample_index];
            }
            selected.s21.push_back(sample);
        }

        traces[key] = std::move(selected);
    }

    return traces;
}

[[nodiscard]] auto compute_ascan(
    const SelectedTrace& trace,
    double start_hz,
    double stop_hz,
    double velocity_mps
) -> AscanResult {
    AscanResult result{};
    if (trace.frequency_hz.size() != trace.s21.size()) {
        return result;
    }

    std::vector<double> frequency_hz{};
    std::vector<std::complex<double>> s21{};
    frequency_hz.reserve(trace.frequency_hz.size());
    s21.reserve(trace.s21.size());

    const double low_hz = std::min(start_hz, stop_hz);
    const double high_hz = std::max(start_hz, stop_hz);
    for (std::size_t index = 0U; index < trace.frequency_hz.size(); ++index) {
        const double frequency_value = trace.frequency_hz[index];
        if (frequency_value < low_hz || frequency_value > high_hz) {
            continue;
        }
        frequency_hz.push_back(frequency_value);
        s21.push_back(trace.s21[index]);
    }

    if (frequency_hz.size() < 2U) {
        return result;
    }

    const std::size_t point_count = frequency_hz.size();
    const double df = (frequency_hz.back() - frequency_hz.front()) / static_cast<double>(point_count - 1U);
    if (!(df > 0.0)) {
        return result;
    }

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

    const auto start_index = static_cast<std::size_t>(start_bin);
    const std::size_t min_fft_len = 2U * (start_index + point_count - 1U);
    const std::size_t fft_len = next_power_of_two(min_fft_len);
    if (fft_len < min_fft_len || start_index > fft_len || point_count > (fft_len - start_index)) {
        return result;
    }

    std::vector<std::complex<double>> spectrum(fft_len, std::complex<double>(0.0, 0.0));
    for (std::size_t index = 0U; index < point_count; ++index) {
        const double window = point_count > 1U
            ? 0.5 - (0.5 * std::cos((2.0 * kPi * static_cast<double>(index)) / static_cast<double>(point_count - 1U)))
            : 1.0;
        spectrum[start_index + index] = s21[index] * window;
    }

    fft_inplace(spectrum, true);

    result.bandwidth_hz = frequency_hz.back() - frequency_hz.front();
    const double dt = 1.0 / (static_cast<double>(fft_len) * df);
    result.time_s.resize(fft_len, 0.0);
    result.depth_m.resize(fft_len, 0.0);
    result.amplitude.resize(fft_len, 0.0);
    for (std::size_t index = 0U; index < fft_len; ++index) {
        result.time_s[index] = static_cast<double>(index) * dt;
        result.depth_m[index] = result.time_s[index] * velocity_mps / 2.0;
        result.amplitude[index] = std::abs(spectrum[index]);
    }

    return result;
}

[[nodiscard]] auto build_grid(
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx,
    double max_depth_m
) -> GridDefinition {
    GridDefinition grid{};
    if (x_tx.empty() || x_rx.empty() || !(max_depth_m > kAccumulatorZMinM)) {
        return grid;
    }

    const auto [tx_min_it, tx_max_it] = std::minmax_element(x_tx.begin(), x_tx.end());
    const auto [rx_min_it, rx_max_it] = std::minmax_element(x_rx.begin(), x_rx.end());
    const double x_min = std::min(*tx_min_it, *rx_min_it) - kAccumulatorXMarginM;
    const double x_max = std::max(*tx_max_it, *rx_max_it) + kAccumulatorXMarginM;

    grid.x_grid = build_axis(x_min, x_max, kGridWidth);
    grid.z_grid = build_axis(kAccumulatorZMinM, max_depth_m, kGridHeight);

    const std::size_t cell_count = grid.x_grid.size() * grid.z_grid.size();
    grid.tx_distance_grids.assign(x_tx.size(), std::vector<double>(cell_count, 0.0));
    grid.rx_distance_grids.assign(x_rx.size(), std::vector<double>(cell_count, 0.0));

    for (std::size_t tx_index = 0U; tx_index < x_tx.size(); ++tx_index) {
        for (std::size_t row = 0U; row < grid.z_grid.size(); ++row) {
            const double z_value = grid.z_grid[row];
            for (std::size_t col = 0U; col < grid.x_grid.size(); ++col) {
                const double x_value = grid.x_grid[col];
                const auto cell_index = (row * grid.x_grid.size()) + col;
                grid.tx_distance_grids[tx_index][cell_index] =
                    std::sqrt(std::pow(x_value - x_tx[tx_index], 2.0) + std::pow(z_value, 2.0));
            }
        }
    }

    for (std::size_t rx_index = 0U; rx_index < x_rx.size(); ++rx_index) {
        for (std::size_t row = 0U; row < grid.z_grid.size(); ++row) {
            const double z_value = grid.z_grid[row];
            for (std::size_t col = 0U; col < grid.x_grid.size(); ++col) {
                const double x_value = grid.x_grid[col];
                const auto cell_index = (row * grid.x_grid.size()) + col;
                grid.rx_distance_grids[rx_index][cell_index] =
                    std::sqrt(std::pow(x_value - x_rx[rx_index], 2.0) + std::pow(z_value, 2.0));
            }
        }
    }

    return grid;
}

[[nodiscard]] auto is_excluded(double z_value, const std::vector<std::pair<double, double>>& ranges) -> bool {
    for (const auto& [low, high] : ranges) {
        if (z_value >= low && z_value <= high) {
            return true;
        }
    }
    return false;
}

[[nodiscard]] auto build_accumulator(
    const GridDefinition& grid,
    const std::unordered_map<PairKey, std::vector<PeakRecord>>& peaks_by_pair,
    const std::vector<std::pair<double, double>>& exclude_ranges,
    double velocity_mps,
    double shell_sigma_m,
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx
) -> std::vector<double> {
    const std::size_t width = grid.x_grid.size();
    const std::size_t height = grid.z_grid.size();
    std::vector<double> accumulator(width * height, 0.0);
    if (!(shell_sigma_m > 0.0)) {
        return accumulator;
    }

    for (std::size_t tx_index = 0U; tx_index < x_tx.size(); ++tx_index) {
        for (std::size_t rx_index = 0U; rx_index < x_rx.size(); ++rx_index) {
            const auto peak_it = peaks_by_pair.find(make_pair_key(static_cast<std::uint32_t>(tx_index), static_cast<std::uint32_t>(rx_index)));
            if (peak_it == peaks_by_pair.end()) {
                continue;
            }

            const auto& tx_grid = grid.tx_distance_grids[tx_index];
            const auto& rx_grid = grid.rx_distance_grids[rx_index];
            for (const auto& peak : peak_it->second) {
                if (is_excluded(peak.z_app, exclude_ranges)) {
                    continue;
                }

                const double range_total = velocity_mps * peak.tau;
                for (std::size_t cell_index = 0U; cell_index < accumulator.size(); ++cell_index) {
                    const double residual = tx_grid[cell_index] + rx_grid[cell_index] - range_total;
                    const double shell = std::exp(-0.5 * std::pow(residual / shell_sigma_m, 2.0));
                    accumulator[cell_index] += shell * peak.snr_comp;
                }
            }
        }
    }

    return accumulator;
}

[[nodiscard]] auto count_agreeing_ellipses(
    double x_est,
    double z_est,
    const std::unordered_map<PairKey, std::vector<PeakRecord>>& peaks_by_pair,
    const std::vector<std::pair<double, double>>& exclude_ranges,
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx,
    double velocity_mps,
    double shell_sigma_m
) -> double {
    std::size_t count = 0U;
    for (std::size_t tx_index = 0U; tx_index < x_tx.size(); ++tx_index) {
        for (std::size_t rx_index = 0U; rx_index < x_rx.size(); ++rx_index) {
            const auto peak_it = peaks_by_pair.find(make_pair_key(static_cast<std::uint32_t>(tx_index), static_cast<std::uint32_t>(rx_index)));
            if (peak_it == peaks_by_pair.end()) {
                continue;
            }

            for (const auto& peak : peak_it->second) {
                if (is_excluded(peak.z_app, exclude_ranges)) {
                    continue;
                }

                const double rt = std::sqrt(std::pow(x_est - x_tx[tx_index], 2.0) + std::pow(z_est, 2.0));
                const double rr = std::sqrt(std::pow(x_est - x_rx[rx_index], 2.0) + std::pow(z_est, 2.0));
                if (std::abs((rt + rr) - (velocity_mps * peak.tau)) < shell_sigma_m * 6.0) {
                    count += 1U;
                    break;
                }
            }
        }
    }

    return static_cast<double>(count);
}

[[nodiscard]] auto find_centroid(
    const std::vector<double>& values,
    std::size_t width,
    std::size_t height,
    std::size_t iz,
    std::size_t ix,
    std::size_t radius_z,
    std::size_t radius_x,
    const std::vector<double>& x_grid,
    const std::vector<double>& z_grid
) -> std::pair<double, double> {
    if (values.empty()) {
        return {x_grid[ix], z_grid[iz]};
    }

    const auto row_start = iz > radius_z ? iz - radius_z : 0U;
    const auto row_stop = std::min(height, iz + radius_z + 1U);
    const auto col_start = ix > radius_x ? ix - radius_x : 0U;
    const auto col_stop = std::min(width, ix + radius_x + 1U);

    double weight_sum = 0.0;
    double row_weighted_sum = 0.0;
    double col_weighted_sum = 0.0;
    for (std::size_t row = row_start; row < row_stop; ++row) {
        for (std::size_t col = col_start; col < col_stop; ++col) {
            const double weight = values[(row * width) + col];
            weight_sum += weight;
            row_weighted_sum += static_cast<double>(row) * weight;
            col_weighted_sum += static_cast<double>(col) * weight;
        }
    }

    if (!(weight_sum > 0.0)) {
        return {x_grid[ix], z_grid[iz]};
    }

    const auto centroid_row = clamp_index(static_cast<std::ptrdiff_t>(std::llround(row_weighted_sum / weight_sum)), height);
    const auto centroid_col = clamp_index(static_cast<std::ptrdiff_t>(std::llround(col_weighted_sum / weight_sum)), width);
    return {x_grid[centroid_col], z_grid[centroid_row]};
}

[[nodiscard]] auto flatten_table(
    const std::vector<std::vector<float>>& rows,
    std::uint32_t column_count
) -> std::vector<float> {
    std::vector<float> values{};
    values.reserve(rows.size() * column_count);
    for (const auto& row : rows) {
        values.insert(values.end(), row.begin(), row.end());
    }
    return values;
}

[[nodiscard]] auto build_table_payload(
    const std::string& processing_name,
    const std::vector<std::vector<float>>& rows,
    std::uint32_t column_count
) -> ipc::ResultPayload {
    ipc::ResultPayload payload{};
    payload.processing_name = processing_name;
    payload.kind = ipc::ResultKind::TableF32;
    payload.table_columns = column_count;
    payload.table_values = flatten_table(rows, column_count);
    return payload;
}

[[nodiscard]] auto build_image_payload(
    const std::string& processing_name,
    const std::vector<double>& x_axis,
    const std::vector<double>& y_axis,
    const std::vector<double>& values
) -> ipc::ResultPayload {
    ipc::ResultPayload payload{};
    payload.processing_name = processing_name;
    payload.kind = ipc::ResultKind::ImageF32;
    payload.image_x_axis.reserve(x_axis.size());
    for (const auto value : x_axis) {
        payload.image_x_axis.push_back(static_cast<float>(value));
    }
    payload.image_y_axis.reserve(y_axis.size());
    for (const auto value : y_axis) {
        payload.image_y_axis.push_back(static_cast<float>(value));
    }
    payload.image_values.reserve(values.size());
    for (const auto value : values) {
        payload.image_values.push_back(static_cast<float>(value));
    }
    return payload;
}

[[nodiscard]] auto accumulator_max(const std::vector<double>& values) -> double {
    if (values.empty()) {
        return 0.0;
    }
    return *std::max_element(values.begin(), values.end());
}

[[nodiscard]] auto clean_find_points(
    const GridDefinition& grid,
    const std::unordered_map<PairKey, std::vector<PeakRecord>>& peaks_by_pair,
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx,
    double velocity_mps,
    double shell_sigma_m
) -> std::pair<std::vector<PointRecord>, std::vector<double>> {
    std::vector<PointRecord> found{};
    const auto accumulator = build_accumulator(grid, peaks_by_pair, {}, velocity_mps, shell_sigma_m, x_tx, x_rx);
    const double initial_max = accumulator_max(accumulator);
    if (!(initial_max > 0.0) || grid.x_grid.size() < 2U || grid.z_grid.size() < 2U) {
        return {found, gaussian_filter_2d(accumulator, grid.x_grid.size(), grid.z_grid.size(), kGaussianSigma)};
    }

    const double dx = grid.x_grid[1] - grid.x_grid[0];
    const double dz = grid.z_grid[1] - grid.z_grid[0];
    const auto radius_x = static_cast<std::size_t>(std::max(1.0, std::round(kCleanSuppressRadiusM / std::max(dx, 1e-6))));
    const auto radius_z = static_cast<std::size_t>(std::max(1.0, std::round(kCleanSuppressRadiusM / std::max(dz, 1e-6))));

    std::vector<std::pair<double, double>> excluded_ranges{};
    for (std::size_t step = 0U; step < kMaxObjects; ++step) {
        const auto current = build_accumulator(grid, peaks_by_pair, excluded_ranges, velocity_mps, shell_sigma_m, x_tx, x_rx);
        const auto smoothed = gaussian_filter_2d(current, grid.x_grid.size(), grid.z_grid.size(), kGaussianSigma);
        const double smoothed_max = accumulator_max(smoothed);
        if (!(smoothed_max > (kCleanThresholdFrac * initial_max))) {
            break;
        }

        const auto max_it = std::max_element(smoothed.begin(), smoothed.end());
        const auto max_index = static_cast<std::size_t>(std::distance(smoothed.begin(), max_it));
        const auto iz = max_index / grid.x_grid.size();
        const auto ix = max_index % grid.x_grid.size();
        const auto [x_est, z_est] = find_centroid(
            smoothed,
            grid.x_grid.size(),
            grid.z_grid.size(),
            iz,
            ix,
            radius_z,
            radius_x,
            grid.x_grid,
            grid.z_grid
        );

        PointRecord point{};
        point.x_m = x_est;
        point.z_m = z_est;
        point.score = count_agreeing_ellipses(x_est, z_est, peaks_by_pair, excluded_ranges, x_tx, x_rx, velocity_mps, shell_sigma_m);
        found.push_back(point);

        std::vector<double> matched_depths{};
        for (std::size_t tx_index = 0U; tx_index < x_tx.size(); ++tx_index) {
            for (std::size_t rx_index = 0U; rx_index < x_rx.size(); ++rx_index) {
                const auto peak_it = peaks_by_pair.find(make_pair_key(static_cast<std::uint32_t>(tx_index), static_cast<std::uint32_t>(rx_index)));
                if (peak_it == peaks_by_pair.end()) {
                    continue;
                }
                for (const auto& peak : peak_it->second) {
                    if (is_excluded(peak.z_app, excluded_ranges)) {
                        continue;
                    }
                    const double rt = std::sqrt(std::pow(x_est - x_tx[tx_index], 2.0) + std::pow(z_est, 2.0));
                    const double rr = std::sqrt(std::pow(x_est - x_rx[rx_index], 2.0) + std::pow(z_est, 2.0));
                    if (std::abs((rt + rr) - (velocity_mps * peak.tau)) < shell_sigma_m * 3.0) {
                        matched_depths.push_back(peak.z_app);
                    }
                }
            }
        }

        if (!matched_depths.empty()) {
            const auto [min_it, max_it_depth] = std::minmax_element(matched_depths.begin(), matched_depths.end());
            excluded_ranges.emplace_back(*min_it - shell_sigma_m, *max_it_depth + shell_sigma_m);
        }
    }

    return {found, gaussian_filter_2d(accumulator, grid.x_grid.size(), grid.z_grid.size(), kGaussianSigma)};
}

[[nodiscard]] auto extended_find_regions(
    const GridDefinition& grid,
    const std::unordered_map<PairKey, std::vector<PeakRecord>>& peaks_by_pair,
    const std::vector<double>& x_tx,
    const std::vector<double>& x_rx,
    double velocity_mps,
    double shell_sigma_m
) -> std::pair<std::vector<RegionRecord>, std::vector<double>> {
    std::vector<RegionRecord> regions{};
    const auto accumulator = build_accumulator(grid, peaks_by_pair, {}, velocity_mps, shell_sigma_m, x_tx, x_rx);
    const auto smoothed = gaussian_filter_2d(accumulator, grid.x_grid.size(), grid.z_grid.size(), kGaussianSigma);
    const double smoothed_max = accumulator_max(smoothed);
    if (!(smoothed_max > 0.0) || grid.x_grid.size() < 2U || grid.z_grid.size() < 2U) {
        return {regions, smoothed};
    }

    const double dx_cm = (grid.x_grid[1] - grid.x_grid[0]) * 100.0;
    const double dz_cm = (grid.z_grid[1] - grid.z_grid[0]) * 100.0;
    const auto min_pixels = static_cast<std::size_t>(std::max(1.0, std::floor(kExtendedMinAreaCm2 / std::max(dx_cm * dz_cm, 1e-6))));

    const std::size_t width = grid.x_grid.size();
    const std::size_t height = grid.z_grid.size();
    const double threshold = kExtendedThresholdFrac * smoothed_max;
    std::vector<std::uint8_t> visited(width * height, 0U);

    for (std::size_t row = 0U; row < height; ++row) {
        for (std::size_t col = 0U; col < width; ++col) {
            const auto start_index = (row * width) + col;
            if (visited[start_index] != 0U || smoothed[start_index] <= threshold) {
                continue;
            }

            std::vector<std::size_t> stack{start_index};
            std::vector<std::size_t> component{};
            visited[start_index] = 1U;

            while (!stack.empty()) {
                const auto cell_index = stack.back();
                stack.pop_back();
                component.push_back(cell_index);

                const auto cell_row = cell_index / width;
                const auto cell_col = cell_index % width;
                const std::pair<std::ptrdiff_t, std::ptrdiff_t> offsets[] = {
                    {-1, 0},
                    {1, 0},
                    {0, -1},
                    {0, 1},
                };

                for (const auto& [row_offset, col_offset] : offsets) {
                    const auto next_row = static_cast<std::ptrdiff_t>(cell_row) + row_offset;
                    const auto next_col = static_cast<std::ptrdiff_t>(cell_col) + col_offset;
                    if (next_row < 0 || next_col < 0) {
                        continue;
                    }
                    if (next_row >= static_cast<std::ptrdiff_t>(height) || next_col >= static_cast<std::ptrdiff_t>(width)) {
                        continue;
                    }
                    const auto next_index =
                        (static_cast<std::size_t>(next_row) * width) + static_cast<std::size_t>(next_col);
                    if (visited[next_index] != 0U || smoothed[next_index] <= threshold) {
                        continue;
                    }
                    visited[next_index] = 1U;
                    stack.push_back(next_index);
                }
            }

            if (component.size() < min_pixels) {
                continue;
            }

            RegionRecord region{};
            region.mask.assign(width * height, 0.0F);
            double weight_sum = 0.0;
            double x_weight_sum = 0.0;
            double z_weight_sum = 0.0;
            for (const auto cell_index : component) {
                const auto cell_row = cell_index / width;
                const auto cell_col = cell_index % width;
                const double weight = smoothed[cell_index];
                weight_sum += weight;
                x_weight_sum += grid.x_grid[cell_col] * weight;
                z_weight_sum += grid.z_grid[cell_row] * weight;
                region.mask[cell_index] = 1.0F;
            }

            if (!(weight_sum > 0.0)) {
                continue;
            }

            region.x_m = x_weight_sum / weight_sum;
            region.z_m = z_weight_sum / weight_sum;
            region.score = count_agreeing_ellipses(region.x_m, region.z_m, peaks_by_pair, {}, x_tx, x_rx, velocity_mps, shell_sigma_m);
            region.pixel_count = static_cast<double>(component.size());
            regions.push_back(std::move(region));
        }
    }

    return {regions, smoothed};
}

}  // namespace

auto GprProcessor::name() const -> std::string {
    return "gpr";
}

auto GprProcessor::process_collection(
    const config::RunConfig& run_config,
    const ipc::PreprocessedCollection& collection,
    std::span<const ipc::PreprocessedCollection> previous_collections,
    const ProcessingLiveConfig& live_config
) -> ipc::ResultCollection {
    ipc::ResultCollection results{};
    results.collection_id = collection.collection_id;
    results.monotonic_ns = collection.monotonic_ns;

    const auto selection = build_geometry_selection(run_config, live_config);
    if (selection.input_positions.empty() || selection.output_positions.empty()) {
        return results;
    }

    const auto background_mean = build_background_mean(previous_collections, selection, live_config);
    const auto traces_by_pair = collect_selected_traces(collection, selection, background_mean);
    if (traces_by_pair.empty()) {
        return results;
    }

    const double velocity_mps =
        kSpeedOfLightMetersPerSec / std::sqrt(std::max(1e-6, static_cast<double>(run_config.gpr.relative_permittivity)));
    const double start_hz = static_cast<double>(live_config.gpr_start_freq_mhz) * 1'000'000.0;
    const double stop_hz = static_cast<double>(live_config.gpr_stop_freq_mhz) * 1'000'000.0;

    std::unordered_map<PairKey, AscanResult> ascans_by_pair{};
    double bandwidth_hz = 0.0;
    for (const auto& [key, trace] : traces_by_pair) {
        auto ascan = compute_ascan(trace, start_hz, stop_hz, velocity_mps);
        if (ascan.amplitude.empty() || !(ascan.bandwidth_hz > 0.0)) {
            continue;
        }
        bandwidth_hz = std::max(bandwidth_hz, ascan.bandwidth_hz);
        ascans_by_pair.emplace(key, std::move(ascan));
    }
    if (ascans_by_pair.empty() || !(bandwidth_hz > 0.0)) {
        return results;
    }

    const double max_depth_m = static_cast<double>(live_config.gpr_max_depth_m);
    const auto grid = build_grid(selection.x_tx, selection.x_rx, max_depth_m);
    if (grid.x_grid.empty() || grid.z_grid.empty()) {
        return results;
    }

    const double shell_sigma_m = velocity_mps / bandwidth_hz * 0.5;
    std::unordered_map<PairKey, std::vector<PeakRecord>> peaks_by_pair{};
    for (std::size_t tx_index = 0U; tx_index < selection.x_tx.size(); ++tx_index) {
        for (std::size_t rx_index = 0U; rx_index < selection.x_rx.size(); ++rx_index) {
            const auto key = make_pair_key(static_cast<std::uint32_t>(tx_index), static_cast<std::uint32_t>(rx_index));
            const auto ascan_it = ascans_by_pair.find(key);
            if (ascan_it == ascans_by_pair.end()) {
                continue;
            }

            const auto& ascan = ascan_it->second;
            if (ascan.depth_m.size() < 3U || ascan.amplitude.size() < 3U) {
                continue;
            }

            const double noise = median_copy(ascan.amplitude);
            const auto min_index = lower_bound_index(ascan.depth_m, static_cast<double>(live_config.gpr_min_depth_m));
            const auto max_index = lower_bound_index(ascan.depth_m, static_cast<double>(live_config.gpr_max_depth_m));
            if (max_index <= min_index + 2U) {
                continue;
            }

            const double z_step = std::max(ascan.depth_m[1] - ascan.depth_m[0], 1e-6);
            const std::size_t min_distance = static_cast<std::size_t>(std::max(
                4.0,
                std::floor(((velocity_mps / (2.0 * bandwidth_hz)) / z_step) * 0.7)
            ));
            const auto peak_indices = find_peak_indices(ascan.amplitude, min_index, max_index, noise * kSnrThresh, min_distance);
            auto& peaks = peaks_by_pair[key];
            peaks.reserve(peak_indices.size());
            for (const auto peak_index : peak_indices) {
                const double z_app = ascan.depth_m[peak_index];
                const double snr_raw = ascan.amplitude[peak_index] / std::max(noise, 1e-12);
                const double attenuation = attenuation_at_depth(tx_index, rx_index, z_app, selection.x_tx, selection.x_rx);
                const double attenuation_norm =
                    attenuation / attenuation_at_depth(tx_index, rx_index, 2.0, selection.x_tx, selection.x_rx);
                const double snr_comp = std::min(
                    snr_raw / (std::pow(attenuation_norm, static_cast<double>(live_config.gpr_comp_power)) + 1e-12),
                    kSnrCompMax
                );
                peaks.push_back(PeakRecord{
                    .z_app = z_app,
                    .tau = ascan.time_s[peak_index],
                    .snr_raw = snr_raw,
                    .snr_comp = snr_comp,
                });
            }
        }
    }

    if (peaks_by_pair.empty()) {
        return results;
    }

    if (run_config.gpr.mode == "extended") {
        const auto [regions, smoothed_accumulator] = extended_find_regions(
            grid,
            peaks_by_pair,
            selection.x_tx,
            selection.x_rx,
            velocity_mps,
            shell_sigma_m
        );
        results.collection_payloads.push_back(
            build_image_payload("gpr_accumulator", grid.x_grid, grid.z_grid, smoothed_accumulator)
        );

        std::vector<std::vector<float>> region_rows{};
        region_rows.reserve(regions.size());
        for (std::size_t index = 0U; index < regions.size(); ++index) {
            const auto& region = regions[index];
            region_rows.push_back(
                {
                    static_cast<float>(region.x_m),
                    static_cast<float>(region.z_m),
                    static_cast<float>(region.score),
                    static_cast<float>(region.pixel_count),
                }
            );
            results.collection_payloads.push_back(
                build_image_payload(
                    "gpr_region_mask_" + std::to_string(index),
                    grid.x_grid,
                    grid.z_grid,
                    std::vector<double>(region.mask.begin(), region.mask.end())
                )
            );
        }
        results.collection_payloads.push_back(build_table_payload("gpr_region_centers", region_rows, 4U));
        return results;
    }

    const auto [points, smoothed_accumulator] = clean_find_points(
        grid,
        peaks_by_pair,
        selection.x_tx,
        selection.x_rx,
        velocity_mps,
        shell_sigma_m
    );
    results.collection_payloads.push_back(build_image_payload("gpr_accumulator", grid.x_grid, grid.z_grid, smoothed_accumulator));

    std::vector<std::vector<float>> point_rows{};
    point_rows.reserve(points.size());
    for (const auto& point : points) {
        point_rows.push_back(
            {
                static_cast<float>(point.x_m),
                static_cast<float>(point.z_m),
                static_cast<float>(point.score),
            }
        );
    }
    results.collection_payloads.push_back(build_table_payload("gpr_points", point_rows, 3U));
    return results;
}

}  // namespace radar::processing
