#include "calibration_master.hpp"

#include <cstddef>
#include <memory>
#include <stdexcept>
#include <vector>

#include <Eigen/Core>

namespace radar::preprocessing {
namespace {

class ThroughCalibrator final : public CalibratorInterface {
  public:
    [[nodiscard]] auto name() const -> std::string override {
        return "through";
    }

    [[nodiscard]] auto apply(
        const std::vector<ipc::Complex32>& measured,
        const std::vector<ipc::Complex32>& calibration
    ) const -> std::vector<ipc::Complex32> override {
        if (measured.size() != calibration.size()) {
            throw std::runtime_error("Through calibration vector size mismatch");
        }

        static_assert(sizeof(ipc::Complex32) == sizeof(float) * 2U, "Complex32 layout must be two contiguous floats");

        constexpr float epsilon = 1e-12F;
        using InterleavedComplexView = Eigen::Matrix<float, Eigen::Dynamic, 2, Eigen::RowMajor>;

        const auto point_count = static_cast<Eigen::Index>(measured.size());
        Eigen::Map<const InterleavedComplexView> measured_view(
            reinterpret_cast<const float*>(measured.data()),
            point_count,
            2
        );
        Eigen::Map<const InterleavedComplexView> calibration_view(
            reinterpret_cast<const float*>(calibration.data()),
            point_count,
            2
        );

        const auto measured_re = measured_view.col(0).array();
        const auto measured_im = measured_view.col(1).array();
        const auto calibration_re = calibration_view.col(0).array();
        const auto calibration_im = calibration_view.col(1).array();

        const Eigen::ArrayXf magnitude_squared = calibration_re.square() + calibration_im.square();
        const auto stable_division_mask = (magnitude_squared > epsilon);

        const Eigen::ArrayXf corrected_re = stable_division_mask.select(
            ((measured_re * calibration_re) + (measured_im * calibration_im)) / magnitude_squared,
            measured_re
        );
        const Eigen::ArrayXf corrected_im = stable_division_mask.select(
            ((measured_im * calibration_re) - (measured_re * calibration_im)) / magnitude_squared,
            measured_im
        );

        std::vector<ipc::Complex32> corrected(measured.size());
        Eigen::Map<InterleavedComplexView> corrected_view(reinterpret_cast<float*>(corrected.data()), point_count, 2);
        corrected_view.col(0) = corrected_re.matrix();
        corrected_view.col(1) = corrected_im.matrix();

        return corrected;
    }
};

}  // namespace

auto make_through_calibrator() -> std::unique_ptr<CalibratorInterface> {
    return std::make_unique<ThroughCalibrator>();
}

}  // namespace radar::preprocessing
