#include "acoustic_analyzer/core/distance_estimator.h"
#include <cmath>
#include <algorithm>
#include <memory>

namespace acoustic {

struct DistanceEstimator::Impl {
    DistanceConfig config;
    float kalman_x_ = 0.0f;
    float kalman_p_ = 1.0f;
    bool kalman_init_ = false;

    explicit Impl(const DistanceConfig& cfg) : config(cfg) {}

    float ComputeSplInternal(const std::vector<float>& audio_signal) {
        if (audio_signal.empty()) return 0.0f;
        double sum = 0.0;
        for (float s : audio_signal) {
            sum += static_cast<double>(s) * s;
        }
        float rms = std::sqrt(static_cast<float>(sum) / audio_signal.size());
        float spl = 20.0f * std::log10(std::max(rms, 1e-10f)) + 94.0f + config.spl_calibration_offset;
        return spl;
    }

    float GetReferenceSpl(const std::string& sound_type) {
        if (sound_type == "gunshot") return config.ref_spl_gunshot;
        if (sound_type == "artillery") return config.ref_spl_artillery;
        if (sound_type == "explosion") return config.ref_spl_explosion;
        if (sound_type == "threat") return config.ref_spl_gunshot; // binary model uses gunshot ref
        // Fallback: use gunshot reference
        return config.ref_spl_gunshot > 0.0f ? config.ref_spl_gunshot : 150.0f;
    }

    float EstimateInternal(float spl, const std::string& sound_type) {
        float L_ref = GetReferenceSpl(sound_type);
        float alpha = config.attenuation_alpha;
        if (alpha < 1e-6f) alpha = 0.6f;
        float d = std::pow(10.0f, (L_ref - spl) / (20.0f * alpha));
        d = std::max(d, 1.0f);
        d = std::min(d, 5000.0f);
        return d;
    }

    float UpdateKalmanInternal(float raw_distance) {
        float Q = config.kalman_process_noise;
        float R = config.kalman_measurement_noise;

        if (!kalman_init_) {
            kalman_x_ = raw_distance;
            kalman_p_ = 1.0f;
            kalman_init_ = true;
        }

        // Predict
        kalman_p_ += Q;
        // Update
        float K = kalman_p_ / (kalman_p_ + R);
        kalman_x_ += K * (raw_distance - kalman_x_);
        kalman_p_ = (1.0f - K) * kalman_p_;
        return kalman_x_;
    }
};

DistanceEstimator::DistanceEstimator(const DistanceConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

DistanceEstimator::~DistanceEstimator() = default;
DistanceEstimator::DistanceEstimator(DistanceEstimator&&) noexcept = default;
DistanceEstimator& DistanceEstimator::operator=(DistanceEstimator&&) noexcept = default;

float DistanceEstimator::ComputeSpl(const std::vector<float>& audio_signal) {
    return impl_->ComputeSplInternal(audio_signal);
}

float DistanceEstimator::Estimate(float spl, const std::string& sound_type) {
    return impl_->EstimateInternal(spl, sound_type);
}

float DistanceEstimator::UpdateKalman(float raw_distance) {
    return impl_->UpdateKalmanInternal(raw_distance);
}

void DistanceEstimator::Reset() {
    impl_->kalman_init_ = false;
    impl_->kalman_p_ = 1.0f;
}

const DistanceConfig& DistanceEstimator::Config() const noexcept {
    return impl_->config;
}

}  // namespace acoustic
