#include "reference_master.hpp"

#include <algorithm>
#include <cmath>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include <Eigen/Core>

#include "calibration_master.hpp"

namespace radar::preprocessing {
namespace {

constexpr float kFrequencyRelativeTolerance = 1e-5F;
constexpr float kFrequencyAbsoluteTolerance = 1e-3F;

[[nodiscard]] auto combo_to_string(const ipc::ComboKey& combo) -> std::string {
    return "input=" + std::to_string(combo.input_pos) + " output=" + std::to_string(combo.output_pos);
}

void validate_channel_trace_layout(
    const ChannelTrace& trace,
    const std::string& trace_label,
    const ipc::ComboKey& combo
) {
    if (trace.frequency_hz.size() != trace.samples.size()) {
        throw std::runtime_error(
            trace_label + " frequency/channel vector size mismatch for combo " + combo_to_string(combo)
        );
    }
}

[[nodiscard]] auto frequency_axes_match(std::span<const float> left, std::span<const float> right) -> bool {
    if (left.size() != right.size()) {
        return false;
    }

    for (std::size_t index = 0; index < left.size(); ++index) {
        const auto delta = std::fabs(left[index] - right[index]);
        const auto scale = std::max(std::fabs(left[index]), std::fabs(right[index]));
        const auto tolerance = std::max(kFrequencyAbsoluteTolerance, scale * kFrequencyRelativeTolerance);
        if (delta > tolerance) {
            return false;
        }
    }
    return true;
}

}  // namespace

void ReferenceMaster::load_s21_reference_bundle(const std::string& path) {
    raw_s21_reference_bundle_.load(path);
    calibrated_s21_references_by_combo_.clear();
}

void ReferenceMaster::load_s11_reference_bundle(const std::string& path) {
    raw_s11_reference_bundle_.load(path);
    calibrated_s11_references_by_combo_.clear();
}

void ReferenceMaster::prepare_calibrated(
    const CalibrationMaster& calibration_master,
    const std::vector<ipc::ComboKey>& combos
) {
    if (!raw_s21_reference_bundle_.is_enabled()) {
        throw std::runtime_error("Reference bundle must be loaded before prepare_calibrated()");
    }

    calibrated_s21_references_by_combo_.clear();
    calibrated_s21_references_by_combo_.reserve(combos.size());
    for (const auto& combo : combos) {
        const auto& raw_reference = raw_s21_reference_bundle_.resolve(combo);
        ChannelTrace calibrated{};
        calibrated.frequency_hz = raw_reference.frequency_hz;
        calibrated.samples = calibration_master.apply_s21(combo, raw_reference.frequency_hz, raw_reference.samples);
        calibrated_s21_references_by_combo_.insert_or_assign(combo, std::move(calibrated));
    }

    calibrated_s11_references_by_combo_.clear();
    if (!raw_s11_reference_bundle_.is_enabled()) {
        return;
    }
    if (!calibration_master.has_s11_calibration()) {
        throw std::runtime_error("S11 reference bundle requires S11 calibration bundle");
    }

    calibrated_s11_references_by_combo_.reserve(combos.size());
    for (const auto& combo : combos) {
        const auto& raw_reference = raw_s11_reference_bundle_.resolve(combo);
        ChannelTrace calibrated{};
        calibrated.frequency_hz = raw_reference.frequency_hz;
        calibrated.samples = calibration_master.apply_s11(combo, raw_reference.frequency_hz, raw_reference.samples);
        calibrated_s11_references_by_combo_.insert_or_assign(combo, std::move(calibrated));
    }
}

void ReferenceMaster::validate_combos(const std::vector<ipc::ComboKey>& combos) const {
    if (calibrated_s21_references_by_combo_.empty()) {
        throw std::runtime_error(
            "Calibrated references are not prepared. Call ReferenceMaster::prepare_calibrated() at startup."
        );
    }

    raw_s21_reference_bundle_.validate_combos(combos);
    for (const auto& combo : combos) {
        if (!calibrated_s21_references_by_combo_.contains(combo)) {
            throw std::runtime_error("Calibrated S21 reference data is missing for combo " + combo_to_string(combo));
        }
    }

    raw_s11_reference_bundle_.validate_combos(combos);
    if (!raw_s11_reference_bundle_.is_enabled()) {
        return;
    }
    if (calibrated_s11_references_by_combo_.empty()) {
        throw std::runtime_error("Calibrated S11 references are not prepared");
    }
    for (const auto& combo : combos) {
        if (!calibrated_s11_references_by_combo_.contains(combo)) {
            throw std::runtime_error("Calibrated S11 reference data is missing for combo " + combo_to_string(combo));
        }
    }
}

auto ReferenceMaster::apply_to_trace(const ipc::SweepTraceBlock& calibrated_trace) const -> ipc::SweepTraceBlock {
    ipc::SweepTraceBlock output{};
    output.combo = calibrated_trace.combo;
    output.frequency_hz = calibrated_trace.frequency_hz;
    output.s21 = apply_s21(calibrated_trace.combo, calibrated_trace.frequency_hz, calibrated_trace.s21);
    output.s11 = apply_s11(calibrated_trace.combo, calibrated_trace.frequency_hz, calibrated_trace.s11);
    return output;
}

auto ReferenceMaster::apply_s21(
    const ipc::ComboKey& combo,
    std::span<const float> frequency_hz,
    const std::vector<ipc::Complex32>& calibrated_s21
) const -> std::vector<ipc::Complex32> {
    const auto found = calibrated_s21_references_by_combo_.find(combo);
    if (found == calibrated_s21_references_by_combo_.end()) {
        throw std::runtime_error("Calibrated S21 reference is missing for combo " + combo_to_string(combo));
    }

    const auto& reference = found->second;
    validate_channel_trace_layout(reference, "Calibrated S21 reference", combo);
    if (calibrated_s21.size() != reference.samples.size()) {
        throw std::runtime_error("S21 reference point count mismatch for combo " + combo_to_string(combo));
    }
    if (!frequency_axes_match(frequency_hz, reference.frequency_hz)) {
        throw std::runtime_error("S21 reference frequency axis mismatch for combo " + combo_to_string(combo));
    }

    std::vector<ipc::Complex32> output(calibrated_s21.size());

    static_assert(sizeof(ipc::Complex32) == sizeof(float) * 2U, "Complex32 layout must be two contiguous floats");
    using InterleavedComplexView = Eigen::Matrix<float, Eigen::Dynamic, 2, Eigen::RowMajor>;
    const auto point_count = static_cast<Eigen::Index>(calibrated_s21.size());

    Eigen::Map<const InterleavedComplexView> calibrated_view(
        reinterpret_cast<const float*>(calibrated_s21.data()),
        point_count,
        2
    );
    Eigen::Map<const InterleavedComplexView> reference_view(
        reinterpret_cast<const float*>(reference.samples.data()),
        point_count,
        2
    );
    Eigen::Map<InterleavedComplexView> output_view(reinterpret_cast<float*>(output.data()), point_count, 2);

    output_view = calibrated_view - reference_view;
    return output;
}

auto ReferenceMaster::apply_s11(
    const ipc::ComboKey& combo,
    std::span<const float> frequency_hz,
    const std::vector<ipc::Complex32>& calibrated_s11
) const -> std::vector<ipc::Complex32> {
    if (!raw_s11_reference_bundle_.is_enabled()) {
        return calibrated_s11;
    }

    const auto found = calibrated_s11_references_by_combo_.find(combo);
    if (found == calibrated_s11_references_by_combo_.end()) {
        throw std::runtime_error("Calibrated S11 reference is missing for combo " + combo_to_string(combo));
    }

    const auto& reference = found->second;
    validate_channel_trace_layout(reference, "Calibrated S11 reference", combo);
    if (calibrated_s11.size() != reference.samples.size()) {
        throw std::runtime_error("S11 reference point count mismatch for combo " + combo_to_string(combo));
    }
    if (!frequency_axes_match(frequency_hz, reference.frequency_hz)) {
        throw std::runtime_error("S11 reference frequency axis mismatch for combo " + combo_to_string(combo));
    }

    std::vector<ipc::Complex32> output(calibrated_s11.size());
    for (std::size_t index = 0; index < calibrated_s11.size(); ++index) {
        output[index].re = calibrated_s11[index].re - reference.samples[index].re;
        output[index].im = calibrated_s11[index].im - reference.samples[index].im;
    }
    return output;
}

}  // namespace radar::preprocessing
