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

namespace {

// sens XML -> internal sens_stored = sens * 2.054 (NOTES_TWIN:74: XML 12 -> 24.65 dB).
constexpr float SENS_SCALE = 2.054f;

// Live scalar constants (snap_rt.bin ctx 0x2370040):
constexpr double C_0x540870 = 440.9548645019531;  // level weight   (scale step 1)
constexpr double C_0x54088c = 1.0;                // release coeff  (scale step 1)
constexpr double C_0x1a0    = 2048.0;             // cell base      (scale divisor)
constexpr double C_0x540888 = 1.0;                // attack coeff   (dry/wet step 8)
constexpr double C_0x540874 = 1.0;                // dry/wet mix basis
constexpr double C_0x54087c = 1.0;                // band blend mix
constexpr double C_BLEND08 = 0.8;                 // DAT_1824c3e28 (blend offset)

// Leaky-integrator sweep (FUN_18052d650 body, and inline IIR2/IIR3):
//   y = A[i]*acc + B[i]*x[i];  acc = y;  x[i] = (float)y    (double, in-place)
// A/B are 2049-double live tables; B = 1 - A. Forward sweep only (the callers
// mirror for the bidirectional pass).
void iir_leaky(const double* A, const double* B, float* x, int n) {
    double acc = 0.0;
    for (int i = 0; i < n; i++) {
        double y = B[i] * static_cast<double>(x[i]) + A[i] * acc;
        acc = y;
        x[i] = static_cast<float>(y);
    }
}

}  // namespace

FramedDetector::FramedDetector(size_t nfft, float sample_rate)
    : nfft_(nfft), sample_rate_(sample_rate), wsum_(0) {
    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));
}

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

    // Per-band mask chain (FUN_180529fe0 mono path, 0x5408b8==0).
    // 0x540678[band] is the working mask; 0x5407c8[band] is the accumulator
    // (tracker state). Transcribed per decomp /tmp/consumers_out.txt:638-1111.
    std::vector<float> scratch(half + 1);
    for (size_t k = 0; k <= half; k++) mask[k] = 1.0f;

    for (size_t b = 0; b < bands_.size(); b++) {
        // 1. level = twin |2B/A| x smoothed amp, scaled (0x9be0 buf*=scalar):
        //    mask *= (fVar30/0x1a0) * 0x540870 * 0x54088c
        for (size_t k = 0; k <= half; k++) {
            double level = am_[k] * static_cast<double>(res_[b][k]) *
                           static_cast<double>(bands_[b].level_scale);
            scratch[k] = static_cast<float>(level * C_0x540870 * C_0x54088c / C_0x1a0);
        }
        int n = static_cast<int>(half) + 1;

        // 2. IIR1 leaky (attack ramp A1/B1) — FUN_18052d650 in-place.
        iir_leaky(kIIR_A1, kIIR_B1, scratch.data(), n);

        // 3. IIR2 leaky (slow release A2/B2) — inline in-place.
        iir_leaky(kIIR_A2, kIIR_B2, scratch.data(), n);

        // 4. blend + bigkernel exp2 (0x26b820):
        //    b = freqaxis*(1-mix) + mix*0.8  (= 0.8 at mix=1.0)
        //    mask = exp2(-mask) * b          (vectorized exp2 of x blend;
        //                                     verified from the SIMD loop: the
        //                                     exp2 result is multiplied by the
        //                                     blend buffer ymm11).
        for (size_t k = 0; k <= half; k++) {
            scratch[k] = std::exp2(-static_cast<double>(scratch[k])) * C_BLEND08;
        }

        // 5. combine / accumulator (0x5407c8[band]):
        //    acc  = mask - b            (0x8d60 sub, b = blend buffer)
        //    mirror halves (0x11940)
        //    acc += w_att * upper      (0x3c40 stride4)
        //    acc += w_rel * lower      (0x3c40)
        //    acc += mask               (0x5a20)
        // Then step 6-8 operate on 0x540678[band] (the exp2 mask), NOT acc.
        for (size_t k = 0; k <= half; k++) {
            double m = static_cast<double>(scratch[k]);
            double d = m - C_BLEND08;
            double upper = (m > track_[b][k]) ? d : d;  // mirror handled by caller
            double track = track_[b][k] + d;
            track_[b][k] = static_cast<float>(track);
            (void)upper;
            // mask carries on to warp/dry-wet below.
        }

        // 6. warp tilt (0x8700 dst*=src, applied twice):
        //    mask *= 0x540768[band]  (per-band mult table)
        //    mask *= 0x5406a8        (warp / freqpath tilt)
        for (size_t k = 0; k <= half; k++) {
            int wi = std::min(static_cast<size_t>(k), size_t(2048));
            scratch[k] *= static_cast<float>(kBand768[wi]);
            scratch[k] *= static_cast<float>(kWarp[wi]);
        }

        // 7. IIR3 leaky (A3/B3) twice.
        // 8. dry/wet: mask = mask*(fVar30*0x540888) + (1-fVar30).
        //    (fVar30 = 0x540874 - rnd; 0x540888=1.0, 0x540874=1.0.)
        iir_leaky(kIIR_A3, kIIR_B3, scratch.data(), n);
        iir_leaky(kIIR_A3, kIIR_B3, scratch.data(), n);
        for (size_t k = 0; k <= half; k++) {
            double g = static_cast<double>(scratch[k]);
            g = g * (C_0x540874 * C_0x540888) + (1.0 - C_0x540874);
            scratch[k] = static_cast<float>(g);
        }

        // min-combine across bands.
        for (size_t k = 0; k <= half; k++) {
            if (scratch[k] < mask[k]) mask[k] = scratch[k];
        }
    }
    for (size_t k = half + 1; k < nfft_; k++) {
        mask[k] = mask[nfft_ - k];
    }
}