#include "framed_model.hpp"
#include "twin.hpp"
#include "freqpath.hpp"
#include "rt_mask_tables.hpp"
#include "rt_weights.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_MIX      = 1.0;                // 0x54087c band blend mix (constructor 0.5? live=1.0)
constexpr double C_0x54087c = C_MIX;
constexpr double C_BLEND08  = 0.8;                // DAT_1824c3e28 (blend offset)
constexpr double C_AXIS     = 1.3;                // 0x540698 freq-axis placeholder (offline = per-band scalar; ref: 8.3-7/(1+exp..) ~1.3)

// 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 live tables; B = 1 - A. Forward sweep only.
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);
    }
}

// 0x11940 copy/mirror: upper half = reversed lower half (Hermitian-style mirror of
// [0,n) into [n, nfft)). In-place on a full-spectrum buffer.
void mirror_half(float* buf, int n, int nfft) {
    for (int i = 0; i < n && (nfft - 1 - i) >= n; i++) buf[nfft - 1 - i] = buf[i];
}

// FUN_180529fe0 PRNG prologue (:515-583). LCG state advances by round offsets;
// fVar30 (scale coeff) = (int)(LUT[s+1]*LUT[s]+0.001). CONSTANTS b5c8/b704/b700
// resolve to 1 (VA-linear dump). Live state 112 => fVar30 == 1.0 deterministically.
double prng_fvar30(int& state) {
    constexpr unsigned M = 0x8000007f;
    unsigned s = static_cast<unsigned>(state & 0xffffffff);
    auto fix = [](unsigned x) {
        x &= M;
        if (x & 0x80000000u) x = ((x - 1) | 0xffffff80u) + 1;
        return x;
    };
    s = fix(s + 0x3cdca);          // line 515
    s = fix(s + 0x140236);         // line 526
    s = fix(s + 0x10d56);          // line 538
    s = fix(s + 0xdf6b6);          // line 549 (s3 -> fVar30)
    unsigned s3 = s;
    int lu = static_cast<int>(s3 & 0x7f);
    float a = static_cast<float>(kPRNGLut[lu]);
    float b = static_cast<float>(kPRNGLut[(lu + 1) & 0x7f]);
    float prod = a * b + 0.001f;
    int fv = static_cast<int>(prod);          // cvttss2si truncation
    state = static_cast<int>(s3);
    if (fv < 1) fv = 1;
    return static_cast<double>(fv);
}

}  // 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();
    f6f8_.assign(half + 1, 0.0f);

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

    // PRNG prologue: advance state once per frame; fVar30 -> scale coefficient.
    double fVar30 = prng_fvar30(prng_state_);

    // Per-band mask chain (FUN_180529fe0 mono path, 0x5408b8==0).
    // Root source: /tmp/consumers_out.txt:638-1111 + NOTES_LEVEL:199-226 (2026-08-19e).
    std::vector<float> scratch(nfft_, 1.0f);
    std::vector<float> diff(nfft_, 0.0f);
    for (size_t k = 0; k <= half; k++) mask[k] = 1.0f;

    for (size_t b = 0; b < bands_.size(); b++) {
        // 1. scale: 0x540678[band] *= (fVar30/0x1a0)·0x540870·0x54088c    (:656)
        //    Level = am / res (res = |2B/A| is minimal at band centre).
        for (size_t k = 0; k <= half; k++) {
            double level = am_[k] / std::max(static_cast<double>(res_[b][k]), 1e-12) *
                           static_cast<double>(bands_[b].level_scale);
            scratch[k] = static_cast<float>(level * (fVar30 / C_0x1a0) *
                                            C_0x540870 * C_0x54088c);
        }
        int n = static_cast<int>(half) + 1;

        // 2. IIR1 leaky into the SHARED 0x5406f8 buffer (:668 FUN_18052d650),
        //    then copy into the band's 0x540678 (0x5160 bridge, :731).
        iir_leaky(kIIR_A1, kIIR_B1, scratch.data(), n);
        for (size_t k = 0; k <= half; k++) f6f8_[k] = scratch[k];
        for (size_t k = 0; k <= half; k++) scratch[k] = f6f8_[k];

        //    IIR2 leaky (slow release A2/B2, states 0x3404f8/0x2c04f8) — applied to
        //    0x540678[band] in-place (:739-818).
        iir_leaky(kIIR_A2, kIIR_B2, scratch.data(), n);

        //    mirror 0x540678[band] lower->upper (0x11940, :817).
        mirror_half(scratch.data(), n, static_cast<int>(nfft_));

        // 3. mono online blend + bigkernel exp2 (0x26b820):
        //    0x5406f8 = 0x540698·(1−mix) + mix·0.8         (0x6e40 + 0x15060, :814-823)
        //    0x540678[band] = exp2(−0x540678)·0x5406f8     (bigkernel, :832)
        for (size_t k = 0; k <= half; k++) {
            double bl = C_AXIS * (1.0 - C_0x54087c) + C_0x54087c * C_BLEND08;
            f6f8_[k] = static_cast<float>(bl);
            scratch[k] = static_cast<float>(std::exp2(-static_cast<double>(scratch[k])) *
                                            static_cast<double>(f6f8_[k]));
        }

        // 4. combine (0x8d60/0x3c40/0x5a20, :843-885):
        //    track[band] = 0x540678 − 0x5406f8        (sub)
        //    += w_att·0x5406f8_upper (stride4)        (0x5406c8 = kRTAtt)
        //    += w_rel·0x5406f8_lower                  (0x5406e8 = kRTRel)
        //    += 0x540678                               (0x5a20)
        for (size_t k = 0; k <= half; k++) {
            double m = static_cast<double>(scratch[k]);
            double bl = static_cast<double>(f6f8_[k]);
            double d = m - bl;
            diff[k] = static_cast<float>(d);
        }
        const float* wA = kRTAtt;
        const float* wR = kRTRel;
        std::vector<float>& tr = track_[b];
        for (size_t k = 0; k <= half; k++) {
            size_t wi = std::min(k, size_t(half));
            double val = static_cast<double>(diff[k]) +
                         static_cast<double>(wA[wi]) * static_cast<double>(f6f8_[k]) +
                         static_cast<double>(wR[wi]) * static_cast<double>(f6f8_[k]) +
                         static_cast<double>(scratch[k]);
            tr[k] = static_cast<float>(val);
        }

        // 5. WARP tilt: 0x540678 *= 0x540768[band] (band mult), then *= 0x5406a8
        //    (0x8700, :938/:947).
        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]);
        }

        // 6. IIR3 leaky twice (states 0x3c0510/0x440510, :970-1110).
        iir_leaky(kIIR_A3, kIIR_B3, scratch.data(), n);
        iir_leaky(kIIR_A3, kIIR_B3, scratch.data(), n);

        // 7. dry/wet: mask = mask·(fVar30·0x540888) + (1−fVar30).
        //    fVar30 = 0x540874 − rnd; 0x540888=1.0, 0x540874=1.0. (identity at live consts)
        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];
    }
}