#include "channel_bundle_support.hpp"

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstddef>
#include <cstdint>
#include <fstream>
#include <iterator>
#include <span>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include "calibrator_interface.hpp"

namespace radar::preprocessing {
namespace {

enum class RawTraceChannel {
    S11,
    S21,
};

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

[[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);
}

[[nodiscard]] auto read_binary_file(const std::string& path, const std::string& bundle_label)
    -> std::vector<std::uint8_t> {
    std::ifstream stream(path, std::ios::binary);
    if (!stream.is_open()) {
        throw std::runtime_error("Failed to open " + bundle_label + " bundle: " + path);
    }

    return std::vector<std::uint8_t>(std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>());
}

void validate_raw_trace_layout(const ipc::SweepTraceBlock& trace, const std::string& trace_label) {
    if (trace.frequency_hz.size() != trace.s11.size()) {
        throw std::runtime_error(
            trace_label + " frequency/S11 vector size mismatch for combo " + combo_to_string(trace.combo)
        );
    }
    if (trace.frequency_hz.size() != trace.s21.size()) {
        throw std::runtime_error(
            trace_label + " frequency/S21 vector size mismatch for combo " + combo_to_string(trace.combo)
        );
    }
}

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;
}

void ensure_frequency_axis_match(
    std::span<const float> expected,
    std::span<const float> actual,
    const std::string& label
) {
    if (!frequency_axes_match(expected, actual)) {
        throw std::runtime_error(label + " frequency axis mismatch");
    }
}

[[nodiscard]] auto to_std_complex(const ipc::Complex32& value) -> std::complex<float> {
    return std::complex<float>(value.re, value.im);
}

[[nodiscard]] auto from_std_complex(const std::complex<float>& value) -> ipc::Complex32 {
    return ipc::Complex32{
        .re = value.real(),
        .im = value.imag(),
    };
}

[[nodiscard]] auto pick_channel_samples(const ipc::SweepTraceBlock& trace, RawTraceChannel channel)
    -> const std::vector<ipc::Complex32>& {
    if (channel == RawTraceChannel::S11) {
        return trace.s11;
    }
    return trace.s21;
}

[[nodiscard]] auto to_channel_trace(
    const ipc::SweepTraceBlock& trace,
    RawTraceChannel channel,
    const std::string& bundle_label
) -> ChannelTrace {
    validate_raw_trace_layout(trace, bundle_label + " trace");

    ChannelTrace channel_trace{};
    channel_trace.frequency_hz = trace.frequency_hz;
    channel_trace.samples = pick_channel_samples(trace, channel);
    validate_channel_trace_layout(channel_trace, bundle_label + " trace", trace.combo);
    return channel_trace;
}

[[nodiscard]] auto load_channel_traces(
    const std::string& path,
    const std::string& bundle_label,
    RawTraceChannel channel
) -> std::unordered_map<ipc::ComboKey, ChannelTrace, ipc::ComboKeyHash> {
    const auto bytes = read_binary_file(path, bundle_label);
    if (bytes.empty()) {
        throw std::runtime_error(bundle_label + " bundle is empty: " + path);
    }

    const auto collection = ipc::deserialize_raw_collection(bytes);
    std::unordered_map<ipc::ComboKey, ChannelTrace, ipc::ComboKeyHash> traces_by_combo{};
    traces_by_combo.reserve(collection.traces.size());
    for (const auto& trace : collection.traces) {
        traces_by_combo.insert_or_assign(trace.combo, to_channel_trace(trace, channel, bundle_label));
    }

    if (traces_by_combo.empty()) {
        throw std::runtime_error(bundle_label + " bundle does not contain traces: " + path);
    }
    return traces_by_combo;
}

void validate_bundle_combos(
    const std::unordered_map<ipc::ComboKey, ChannelTrace, ipc::ComboKeyHash>& traces_by_combo,
    const std::vector<ipc::ComboKey>& combos,
    const std::string& bundle_label
) {
    if (traces_by_combo.empty()) {
        return;
    }

    for (const auto& combo : combos) {
        if (!traces_by_combo.contains(combo)) {
            throw std::runtime_error(bundle_label + " is missing combo " + combo_to_string(combo));
        }
    }
}

template <typename TValue>
auto resolve_required(
    const std::unordered_map<ipc::ComboKey, TValue, ipc::ComboKeyHash>& values_by_combo,
    const ipc::ComboKey& combo,
    const std::string& value_label
) -> const TValue& {
    const auto found = values_by_combo.find(combo);
    if (found == values_by_combo.end()) {
        throw std::runtime_error(value_label + " is missing for combo " + combo_to_string(combo));
    }
    return found->second;
}

void ensure_combo_coverage(
    const std::unordered_map<ipc::ComboKey, ChannelTrace, ipc::ComboKeyHash>& expected,
    const std::unordered_map<ipc::ComboKey, ChannelTrace, ipc::ComboKeyHash>& actual,
    const std::string& expected_label,
    const std::string& actual_label
) {
    for (const auto& entry : expected) {
        const auto& combo = entry.first;
        if (!actual.contains(combo)) {
            throw std::runtime_error(
                actual_label + " is missing combo " + combo_to_string(combo) +
                " required by " + expected_label
            );
        }
    }
}

[[nodiscard]] auto solve_osl_coefficients(
    const ipc::ComboKey& combo,
    const ChannelTrace& open_trace,
    const ChannelTrace& short_trace,
    const ChannelTrace& load_trace
) -> S11CalibrationBundle::Coefficients {
    validate_channel_trace_layout(open_trace, "S11 open calibration", combo);
    validate_channel_trace_layout(short_trace, "S11 short calibration", combo);
    validate_channel_trace_layout(load_trace, "S11 load calibration", combo);
    ensure_frequency_axis_match(open_trace.frequency_hz, short_trace.frequency_hz, "S11 short calibration");
    ensure_frequency_axis_match(open_trace.frequency_hz, load_trace.frequency_hz, "S11 load calibration");

    S11CalibrationBundle::Coefficients coefficients{};
    coefficients.frequency_hz = open_trace.frequency_hz;
    coefficients.directivity.resize(open_trace.frequency_hz.size());
    coefficients.source_match.resize(open_trace.frequency_hz.size());
    coefficients.reflection_tracking.resize(open_trace.frequency_hz.size());

    for (std::size_t index = 0; index < open_trace.frequency_hz.size(); ++index) {
        const auto load = to_std_complex(load_trace.samples[index]);
        const auto open_delta = to_std_complex(open_trace.samples[index]) - load;
        const auto short_delta = to_std_complex(short_trace.samples[index]) - load;
        const auto denominator = open_delta - short_delta;

        std::complex<float> source_match = std::complex<float>(0.0F, 0.0F);
        std::complex<float> reflection_tracking = std::complex<float>(1.0F, 0.0F);
        if (std::norm(denominator) > kComplexMagnitudeSquaredEpsilon) {
            source_match = (open_delta + short_delta) / denominator;
            reflection_tracking = open_delta * (std::complex<float>(1.0F, 0.0F) - source_match);
        }

        coefficients.directivity[index] = load_trace.samples[index];
        coefficients.source_match[index] = from_std_complex(source_match);
        coefficients.reflection_tracking[index] = from_std_complex(reflection_tracking);
    }

    return coefficients;
}

}  // namespace

void S21CalibrationBundle::load(const std::string& path) {
    traces_by_combo_.clear();

    if (path.empty()) {
        throw std::runtime_error("S21 calibration bundle path must not be empty");
    }

    traces_by_combo_ = load_channel_traces(path, "S21 calibration", RawTraceChannel::S21);
}

auto S21CalibrationBundle::is_enabled() const noexcept -> bool {
    return !traces_by_combo_.empty();
}

void S21CalibrationBundle::validate_combos(const std::vector<ipc::ComboKey>& combos) const {
    validate_bundle_combos(traces_by_combo_, combos, "S21 calibration bundle");
}

auto S21CalibrationBundle::apply(
    const ipc::ComboKey& combo,
    std::span<const float> frequency_hz,
    const std::vector<ipc::Complex32>& measured,
    const CalibratorInterface& calibrator
) const -> std::vector<ipc::Complex32> {
    const auto& calibration_trace = resolve(combo);
    ensure_frequency_axis_match(calibration_trace.frequency_hz, frequency_hz, "S21 calibration");
    if (measured.size() != calibration_trace.samples.size()) {
        throw std::runtime_error("S21 calibration vector size mismatch");
    }
    return calibrator.apply(measured, calibration_trace.samples);
}

auto S21CalibrationBundle::resolve(const ipc::ComboKey& combo) const -> const ChannelTrace& {
    return resolve_required(traces_by_combo_, combo, "S21 calibration trace");
}

void S21ReferenceBundle::load(const std::string& path) {
    traces_by_combo_.clear();

    if (path.empty()) {
        throw std::runtime_error("S21 reference bundle path must not be empty");
    }

    traces_by_combo_ = load_channel_traces(path, "S21 reference", RawTraceChannel::S21);
}

auto S21ReferenceBundle::is_enabled() const noexcept -> bool {
    return !traces_by_combo_.empty();
}

void S21ReferenceBundle::validate_combos(const std::vector<ipc::ComboKey>& combos) const {
    validate_bundle_combos(traces_by_combo_, combos, "S21 reference bundle");
}

auto S21ReferenceBundle::resolve(const ipc::ComboKey& combo) const -> const ChannelTrace& {
    return resolve_required(traces_by_combo_, combo, "S21 reference trace");
}

void S11CalibrationBundle::load(
    const std::string& open_path,
    const std::string& short_path,
    const std::string& load_path
) {
    coefficients_by_combo_.clear();

    if (open_path.empty() && short_path.empty() && load_path.empty()) {
        return;
    }
    if (open_path.empty() || short_path.empty() || load_path.empty()) {
        throw std::runtime_error("S11 calibration requires open, short, and load raw bundle paths");
    }

    const auto open_traces = load_channel_traces(open_path, "S11 open calibration", RawTraceChannel::S11);
    const auto short_traces = load_channel_traces(short_path, "S11 short calibration", RawTraceChannel::S11);
    const auto load_traces = load_channel_traces(load_path, "S11 load calibration", RawTraceChannel::S11);

    ensure_combo_coverage(open_traces, short_traces, "S11 open calibration", "S11 short calibration");
    ensure_combo_coverage(open_traces, load_traces, "S11 open calibration", "S11 load calibration");
    ensure_combo_coverage(short_traces, open_traces, "S11 short calibration", "S11 open calibration");
    ensure_combo_coverage(load_traces, open_traces, "S11 load calibration", "S11 open calibration");

    coefficients_by_combo_.reserve(open_traces.size());
    for (const auto& [combo, open_trace] : open_traces) {
        coefficients_by_combo_.insert_or_assign(
            combo,
            solve_osl_coefficients(combo, open_trace, short_traces.at(combo), load_traces.at(combo))
        );
    }
}

auto S11CalibrationBundle::is_enabled() const noexcept -> bool {
    return !coefficients_by_combo_.empty();
}

void S11CalibrationBundle::validate_combos(const std::vector<ipc::ComboKey>& combos) const {
    if (!is_enabled()) {
        return;
    }

    for (const auto& combo : combos) {
        if (!coefficients_by_combo_.contains(combo)) {
            throw std::runtime_error("S11 calibration data is missing for combo " + combo_to_string(combo));
        }
    }
}

auto S11CalibrationBundle::apply(
    const ipc::ComboKey& combo,
    std::span<const float> frequency_hz,
    const std::vector<ipc::Complex32>& measured
) const -> std::vector<ipc::Complex32> {
    if (!is_enabled()) {
        return measured;
    }
    if (frequency_hz.size() != measured.size()) {
        throw std::runtime_error("S11 calibration input frequency/sample size mismatch");
    }

    const auto& coefficients = resolve(combo);
    ensure_frequency_axis_match(coefficients.frequency_hz, frequency_hz, "S11 calibration");
    if (measured.size() != coefficients.directivity.size()) {
        throw std::runtime_error("S11 calibration vector size mismatch");
    }

    std::vector<ipc::Complex32> corrected{};
    corrected.reserve(measured.size());
    for (std::size_t index = 0; index < measured.size(); ++index) {
        const auto numerator = to_std_complex(measured[index]) - to_std_complex(coefficients.directivity[index]);
        const auto denominator =
            to_std_complex(coefficients.reflection_tracking[index]) +
            (to_std_complex(coefficients.source_match[index]) * numerator);

        if (std::norm(denominator) > kComplexMagnitudeSquaredEpsilon) {
            corrected.push_back(from_std_complex(numerator / denominator));
        } else {
            corrected.push_back(from_std_complex(numerator));
        }
    }
    return corrected;
}

auto S11CalibrationBundle::resolve(const ipc::ComboKey& combo) const -> const Coefficients& {
    return resolve_required(coefficients_by_combo_, combo, "S11 calibration coefficients");
}

void S11ReferenceBundle::load(const std::string& path) {
    traces_by_combo_.clear();

    if (path.empty()) {
        return;
    }

    traces_by_combo_ = load_channel_traces(path, "S11 reference", RawTraceChannel::S11);
}

auto S11ReferenceBundle::is_enabled() const noexcept -> bool {
    return !traces_by_combo_.empty();
}

void S11ReferenceBundle::validate_combos(const std::vector<ipc::ComboKey>& combos) const {
    validate_bundle_combos(traces_by_combo_, combos, "S11 reference bundle");
}

auto S11ReferenceBundle::resolve(const ipc::ComboKey& combo) const -> const ChannelTrace& {
    return resolve_required(traces_by_combo_, combo, "S11 reference trace");
}

}  // namespace radar::preprocessing
