#include "detect.hpp"
#include <cmath>
#include <cstring>
#include <algorithm>

Detector::Detector(size_t nfft, float sample_rate)
    : nfft_(nfft), sample_rate_(sample_rate),
      sharpness_(1.0f), selectivity_(0.5f), depth_(0.0f) {
    envelope_.resize(nfft, 0.0f);
    prev_mask_.resize(nfft, 1.0f);
    smooth_buf_.resize(nfft, 0.0f);
}

void Detector::setParams(float sharpness, float selectivity, float depth) {
    sharpness_ = sharpness;
    selectivity_ = selectivity;
    depth_ = depth;
}

void Detector::processFrame(const std::complex<double>* spectrum, float* mask) {
    std::vector<float> mag(nfft_);
    for (size_t i = 0; i < nfft_; i++) {
        mag[i] = static_cast<float>(std::sqrt(
            spectrum[i].real() * spectrum[i].real() +
            spectrum[i].imag() * spectrum[i].imag()));
    }

    const float alpha_up = 0.1f;
    const float alpha_dn = 0.001f;

    for (size_t i = 0; i < nfft_; i++) {
        if (mag[i] > envelope_[i]) {
            envelope_[i] += alpha_up * (mag[i] - envelope_[i]);
        } else {
            envelope_[i] += alpha_dn * (mag[i] - envelope_[i]);
        }
    }

    for (size_t i = 0; i < nfft_; i++) {
        float ratio = 1.0f;
        if (envelope_[i] > 1e-10f) {
            ratio = mag[i] / envelope_[i];
        }

        float threshold = selectivity_;
        float reduction = 0.0f;

        if (ratio > threshold) {
            float excess = (ratio - threshold) / (1.0f - threshold + 1e-10f);
            reduction = depth_ * std::pow(std::min(excess, 1.0f), sharpness_);
        }

        mask[i] = 1.0f - reduction;
    }

    const float smooth_alpha = 0.3f;
    for (size_t i = 0; i < nfft_; i++) {
        mask[i] = prev_mask_[i] + smooth_alpha * (mask[i] - prev_mask_[i]);
        prev_mask_[i] = mask[i];
    }
}