#include "framed_model.hpp"
#include "twin.hpp"
#include "freqpath.hpp"
#include "rt_weights.hpp"
#include <cmath>
#include <cstring>
#include <algorithm>

namespace {

// Mask curve fitted against live capture (snap_rt.bin, mask band0 vs acc 0x5407c8):
//   mask = (1 / (1 + K*acc))^n,  K=9.80, n=0.260  (mse 0.019 over 17 resonant bins)
constexpr double MASK_K = 9.8026;
constexpr double MASK_N = 0.25966;
// sens XML -> internal sens_stored = sens * 2.054 (NOTES_TWIN:74: XML 12 -> 24.65 dB).
// GAIN_band = sqrt(param_5) = 10^(sens_stored/40).
constexpr float SENS_SCALE = 2.054f;
constexpr int RT_WEIGHTS_N = 2049;  // captured attack/release tables size

double mask_lut(double acc) {
    if (acc <= 0.0) return 1.0;
    return std::pow(1.0 / (1.0 + MASK_K * acc), MASK_N);
}

}  // namespace

FramedDetector::FramedDetector(size_t nfft, float sample_rate)
    : nfft_(nfft), sample_rate_(sample_rate), wsum_(0) {
    warp_.resize(nfft, 0.0f);
    am_.resize(nfft / 2 + 1, 0.0f);
}

FramedDetector::~FramedDetector() {}

void FramedDetector::setParams(const std::vector<DetectorBand>& bands) {
    bands_ = bands;
    size_t half = nfft_ / 2;
    res_.clear();
    track_.clear();

    for (const auto& b : bands_) {
        std::vector<float> r(half + 1, 1.0f);
        float sens_lin = std::pow(10.0f, b.sens * SENS_SCALE / 20.0f);  // param_5
        detkernel::twin_coeff c = detkernel::build_twin_coeff(
            static_cast<double>(sample_rate_), static_cast<double>(b.fc),
            static_cast<double>(b.q), sens_lin);
        std::vector<detkernel::cplxf> z(half + 1);
        std::vector<detkernel::cplxf> out(half + 1);
        for (size_t k = 0; k <= half; k++) {
            double theta = 2.0 * M_PI * static_cast<double>(k) / static_cast<double>(nfft_);
            z[k].re = static_cast<float>(std::cos(theta));
            z[k].im = static_cast<float>(std::sin(theta));
        }
        detkernel::twin_apply(c, z.data(), half + 1, out.data());
        for (size_t k = 0; k <= half; k++) {
            r[k] = std::sqrt(out[k].re * out[k].re + out[k].im * out[k].im);
            r[k] = std::max(r[k], 1e-12f);
        }
        res_.push_back(std::move(r));
    }
    track_.assign(bands_.size(), std::vector<float>(half + 1, 0.0f));

    if (warp_[0] == 0.0f && warp_[1] == 0.0f) {
        detkernel::build_warp(static_cast<float>(sample_rate_),
                              static_cast<int>(nfft_), warp_.data());
    }
}

void FramedDetector::processFrame(const std::complex<double>* spectrum, float* mask) {
    size_t half = nfft_ / 2;
    if (wsum_ == 0.0) {
        double s = 0.0;
        for (size_t i = 0; i < nfft_; i++) {
            s += std::sqrt(0.5 * (1.0 - std::cos(2.0 * M_PI * i / (nfft_ - 1))));
        }
        wsum_ = s;
    }

    double tatt = 0.011, trel = 0.08;
    double att = std::exp(-1.0 * (nfft_ / 4) / (tatt * sample_rate_));
    double rel = std::exp(-1.0 * (nfft_ / 4) / (trel * sample_rate_));

    for (size_t k = 0; k <= half; k++) {
        double a_cur = 2.0 * std::abs(spectrum[k]) / wsum_;
        double am = am_[k];
        if (a_cur > am) am = att * am + (1.0 - att) * a_cur;
        else am = rel * am + (1.0 - rel) * a_cur;
        am_[k] = static_cast<float>(am);
    }

    for (size_t k = 0; k <= half; k++) {
        double am = am_[k];
        double gain = 1.0;
        for (size_t b = 0; b < bands_.size(); b++) {
            // Level = twin resonance magnitude x smoothed per-bin amplitude.
            double level = am * static_cast<double>(res_[b][k]) *
                           static_cast<double>(bands_[b].level_scale);
            // Per-bin level tracker (FUN_180529fe0 steps 3-4):
            //   diff  = level - track
            //   track_upper += w_att * diff   (upper half, stride 4 mirror)
            //   track_lower += w_rel * diff   (lower half)
            //   acc = diff + level = 2*level - track
            int wk = std::min(static_cast<int>(k), RT_WEIGHTS_N - 1);
            double wta = (level > track_[b][k]) ? static_cast<double>(kRTAtt[wk])
                                                : static_cast<double>(kRTRel[wk]);
            double diff = level - track_[b][k];
            double track = track_[b][k] + wta * diff;
            track_[b][k] = static_cast<float>(track);
            // Accumulator (FUN_180529fe0 step 5: 0x5407c8 = diff; += level).
            double acc = diff + level;
            double g = mask_lut(acc);
            if (g < gain) gain = g;
        }
        mask[k] = static_cast<float>(gain);
    }
    for (size_t k = half + 1; k < nfft_; k++) {
        mask[k] = mask[nfft_ - k];
    }
}