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

namespace {

constexpr float SENS_SCALE = 2.054f;

constexpr double G_FIT = 0.9963;
constexpr double W_FIT = 0.3335;
constexpr double A_FIT = 0.9807;
constexpr double RP0   = 0.0275;
constexpr double DRP   = 0.2159;

static constexpr double kLX[12] = { -0.75, -0.5012, -0.5, -0.2012, 0.0988, 0.2488,
                                     0.3988, 0.5488, 0.574, 0.61,   0.75,   1.0 };
static constexpr double kLY[12] = { 0.4402, 0.366,  0.4552, 0.459, 0.541, 0.576,
                                     0.608,  0.636,  0.5645, 0.6471, 0.6562, 0.6670 };

static double lut_pchip(double x) {
    int n = 12;
    x = std::min(std::max(x, kLX[0]), kLX[n - 1]);
    double h[12], d[12];
    for (int i = 0; i < n - 1; i++) h[i] = kLX[i + 1] - kLX[i];
    for (int i = 0; i < n - 1; i++) d[i] = (kLY[i + 1] - kLY[i]) / h[i];
    double sl[12], sr[12];
    sl[0] = d[0]; sr[n - 1] = d[n - 2];
    for (int i = 1; i < n - 1; i++) {
        if (d[i - 1] * d[i] <= 0.0) { sl[i] = sr[i - 1] = 0.0; continue; }
        double w1 = 2 * h[i] + h[i - 1], w2 = h[i] + 2 * h[i - 1];
        sl[i] = (w1 + w2) / (w1 / d[i - 1] + w2 / d[i]);
        sr[i - 1] = sl[i];
    }
    int i = std::upper_bound(kLX, kLX + n, x) - kLX - 1;
    i = std::max(0, std::min(i, n - 2));
    double hh = h[i], t = (x - kLX[i]) / hh;
    double t2 = t * t, t3 = t2 * t;
    double h00 = 2 * t3 - 3 * t2 + 1, h10 = t3 - 2 * t2 + t;
    double h01 = -2 * t3 + 3 * t2,   h11 = t3 - t2;
    double y = h00 * kLY[i] + h10 * hh * sr[i] + h01 * kLY[i + 1] + h11 * hh * sl[i + 1];
    return y;
}

static double warp_c(double f) {
    double x = f / 2000.0;
    return 0.87 * 7.942 * x / (7.942 + x);
}

static bool is_internal_grid(size_t nfft, float sample_rate) {
    return nfft == 4096 && std::abs(sample_rate - 48000.0f) < 1.0f;
}

static void process_band_structural(
    const float* am,
    const float* res,
    const DetectorBand& band,
    float* mask_out,
    size_t nfft,
    float sample_rate
) {
    const size_t half = nfft / 2;
    const size_t nbin = half + 1;

    static thread_local std::vector<float> band_level;
    static thread_local std::vector<float> f6f8;
    static thread_local std::vector<double> acc;

    band_level.resize(nfft);
    f6f8.resize(nfft);
    acc.assign(nfft, 0.0);

    constexpr float fVar30 = 1.0f;
    constexpr float scale_factor = 15.0f * 440.95f / 2048.0f;
    constexpr float mix = 1.0f;

    // BandConfig ctx+0x188 (FUN_180563a60 dB-domain LUT): A=min, B=max, gamma
    // Extracted from refs: A=-13.78dB, B=68.29dB, gamma=0.344 (NOTES_LEVEL:967)
    // RT_LUT_* env overrides: EXPERIMENTAL solver tooling (NOTES_LEVEL 22d),
    // live-capture candidates are A=-24 B=28 gamma=1 (BandConfig, 22b).
    float lut_a = -13.78f, lut_b = 68.29f, lut_g = 0.344f, lut_m = 4.2f;
    if (const char* e = getenv("RT_LUT_A")) lut_a = atof(e);
    if (const char* e = getenv("RT_LUT_B")) lut_b = atof(e);
    if (const char* e = getenv("RT_LUT_G")) lut_g = atof(e);
    if (const char* e = getenv("RT_LUT_MULT")) lut_m = atof(e);
    const float LUT_A = lut_a;
    const float LUT_B = lut_b;
    const float LUT_GAMMA = lut_g;
    const float LUT_MULT = lut_m;

    // res^rp term (bridge parity): smooth frequency-dependent floor
    constexpr double RP0 = 0.0275;
    constexpr double DRP = 0.2159;
    double rp = RP0 * std::pow(static_cast<double>(band.q), DRP);

    // RT_LUT_OFF=1: EXPERIMENTAL (NOTES 22f) — skip LUT transform entirely,
    // hypothesis: audio path has NO LUT (FUN_180563a60 was GUI-only, 22b);
    // mask = blend*exp2(-lvl_raw) directly.
    static const int lut_off = getenv("RT_LUT_OFF") ? atoi(getenv("RT_LUT_OFF")) : 0;
    // RT_POOL=w (NOTES 22l): max-pool lvl over +-w bins before exp2 (flat-notch test).
    // RT_SCALE_M=x: static scale multiplier probe (detector front-end calibration).
    static const int pool_w = getenv("RT_POOL") ? atoi(getenv("RT_POOL")) : 0;
    static const double scale_mult = getenv("RT_SCALE_M") ? atof(getenv("RT_SCALE_M")) : 1.0;
    const double scale_factor_x = scale_factor * scale_mult;

    std::vector<float> lvl_in(nbin);
    for (size_t k = 0; k < nbin; k++) {
        double res_k = std::max(static_cast<double>(res[k]), 1e-12);
        lvl_in[k] = static_cast<float>(static_cast<double>(am[k]) / res_k * scale_factor_x);
    }
    if (pool_w > 0 && !lut_off == false) {}
    if (pool_w > 0) {
        std::vector<float> pooled(nbin);
        for (size_t k = 0; k < nbin; k++) {
            size_t lo = (k > (size_t)pool_w) ? k - pool_w : 0;
            size_t hi = std::min(nbin - 1, k + (size_t)pool_w);
            float mx = 0.0f;
            for (size_t j = lo; j <= hi; j++) mx = std::max(mx, lvl_in[j]);
            pooled[k] = mx;
        }
        lvl_in.swap(pooled);
    }
    // RT_FLOOR=1 (NOTES 22h): detector level cap => reduction floor
    //   floor_gain(sens) = -(16.78+sens/3)/6.0174*6.0174 dB  => lvl_cap below
    static const int floor_on = getenv("RT_FLOOR") ? atoi(getenv("RT_FLOOR")) : 0;
    if (floor_on) {
        float cap = (16.78f + band.sens / 3.0f) / 6.0174f;
        for (size_t k = 0; k < nbin; k++) if (lvl_in[k] > cap) lvl_in[k] = cap;
    }

    for (size_t k = 0; k < nbin; k++) {
        double res_k = std::max(static_cast<double>(res[k]), 1e-12);
        double lvl = lvl_in[k];
        if (!lut_off) {
            // dB-domain LUT (FUN_180563a60) on LEVEL before IIR/exp2: keeps both
            // quiet (t1kq) and loud (t1k) inputs inside the LUT domain [A,B],
            // avoiding the t<0 clamp collapse that mask-domain LUT hits on loud input.
            double dB = std::log10(std::max(lvl, 1e-12)) * 20.0;
            double t = (dB - LUT_A) / (LUT_B - LUT_A);
            t = std::min(std::max(t, 0.0), 1.0);
            lvl = std::pow(t, static_cast<double>(LUT_GAMMA)) * LUT_MULT;
        }
        band_level[k] = static_cast<float>(lvl);
    }

    // RT_IIR12 mode (NOTES 22j, EXPERIMENTAL): how IIR1/IIR2 run.
    //   fwd (default/canon): ascending-bin cascade within frame.
    //   bidir: forward+backward passes like IIR3.
    //   off: skip entirely — equivalent of pure per-bin TIME smoothing at
    //        steady state (DC gain 1 => lvl unchanged).
    //   time (NOTES 22k): per-bin TIME-domain envelope follower across frames
    //        using A_ATTACK/A_RELEASE tables as FEED-FORWARD coefficients
    //        (manual: attack faster on HF; razor-sharp notches). State persists.
    static const int iir_mode = getenv("RT_IIR12") ? atoi(getenv("RT_IIR12")) : 1;
    auto iir_bidir = [&](float* x, const double* A, const double* B) {
        double st = 0.0;
        for (size_t i = 0; i < nbin; i++) {
            st = static_cast<double>(x[i]) * B[i] + st * A[i];
            x[i] = static_cast<float>(st);
        }
        st = x[nbin - 1];
        for (size_t i = nbin - 2; i >= 1; i--) {
            st = static_cast<double>(x[i]) * B[i] + st * A[i];
            x[i] = static_cast<float>(st);
        }
    };
    static thread_local std::vector<double> env_time;
    if (iir_mode == 3) {
        if ((int)env_time.size() != (int)nbin) env_time.assign(nbin, 0.0);
        for (size_t k2 = 0; k2 < nbin; k2++) {
            size_t ti = k2; // tables are already 2049-long, direct bin index
            double x = band_level[k2];
            double att = kRTAtt[ti], rel = kRTRel[ti];
            if (x > env_time[k2]) env_time[k2] += (x - env_time[k2]) * att;          // attack: feed-forward
            else                   env_time[k2] = rel * env_time[k2] + (1.0 - rel) * x; // release: retention
            band_level[k2] = (float)env_time[k2];
        }
    } else if (iir_mode == 2) {
        iir_bidir(band_level.data(), kIIR_A1, kIIR_B1);
        std::copy(band_level.begin(), band_level.begin() + nbin, f6f8.begin());
        iir_bidir(band_level.data(), kIIR_A2, kIIR_B2);
    } else if (iir_mode == 1) {
        fn529fe0::iir1(band_level.data(), kIIR_A1, kIIR_B1, nbin, 0.0);
        std::copy(band_level.begin(), band_level.begin() + nbin, f6f8.begin());
        fn529fe0::iir1(band_level.data(), kIIR_A2, kIIR_B2, nbin, 0.0);
    }

    // RT_LVL_CAP: EXPERIMENTAL detector-level cap (NOTES 22f/22g/22h) — the real
    // plugin's reduction floors at blend*ln10/20 (sens12/mix100), implying a cap
    // on post-IIR level. Opt-in; default off (canon untouched).
    static const float lvl_cap = getenv("RT_LVL_CAP") ? atof(getenv("RT_LVL_CAP")) : 1e9f;
    for (size_t k = 0; k < nbin; k++) {
        if (band_level[k] > lvl_cap) band_level[k] = lvl_cap;
    }

    for (size_t k = 0; k < half; k++) {
        band_level[nfft - 1 - k] = band_level[k];
    }

    for (size_t k = 0; k < nfft; k++) {
        f6f8[k] = 1.0f * (1.0f - mix) + mix * 0.8f;
    }

    for (size_t k = 0; k < nfft; k++) {
        double mm = std::exp2(-static_cast<double>(band_level[k]));
        static const int noblend = getenv("RT_NOBLEND") ? atoi(getenv("RT_NOBLEND")) : 0;
        if (!noblend) mm *= f6f8[k];
        mask_out[k] = static_cast<float>(mm);
    }

    // RT_DUMP_BIN debug: capture pre-warp mask (opt-in, no cost when unset).
    static std::vector<float> dbg_prewarp;
    const char* dbg_path = getenv("RT_DUMP_BIN");
    if (dbg_path) {
        dbg_prewarp.assign(mask_out, mask_out + nbin);
    }

    fn529fe0::combine_acc(acc.data(), band_level.data(), f6f8.data(),
                          kRTAtt, kRTRel, nfft);

    // RT_NOWARP=1 (NOTES 22j, EXPERIMENTAL): skip warp/W attenuation — white-noise
    // probe shows the real plugin passes broadband content at unity, so the warp
    // term cannot be a blanket output multiplier.
    static const int nowarp = getenv("RT_NOWARP") ? atoi(getenv("RT_NOWARP")) : 0;
    if (!nowarp) {
        for (size_t k = 0; k < nfft; k++) {
            size_t idx = (k < nbin) ? k : (nfft - 1 - k);
            double res_k = std::max(static_cast<double>(res[idx]), 1e-12);
            mask_out[k] *= kBand768[idx] * kWarp[idx] * std::pow(res_k, rp);
        }
    }

    // Step 9 (NOTES_LEVEL:830 + consumers_out.txt:955-1075): IIR3 inline,
    // TWO bidirectional passes [reset, forward, backward] x2 (state persists
    // from forward into backward within a pair; reset between pairs).
    // y = B3[i]*x[i] + A3[i]*state (decomp operand order verified).
    static const int no_iir3 = getenv("RT_NOIIR3") ? atoi(getenv("RT_NOIIR3")) : 0;
    for (int pass = 0; pass < 2 && !no_iir3; pass++) {
        double st = 0.0;
        for (size_t i = 0; i < nbin; i++) {
            double y = static_cast<double>(mask_out[i]) * kIIR_B3[i] + st * kIIR_A3[i];
            st = y;
            mask_out[i] = static_cast<float>(y);
        }
        for (size_t i = nbin - 2; i >= 1; i--) {
            double y = static_cast<double>(mask_out[i]) * kIIR_B3[i] + st * kIIR_A3[i];
            st = y;
            mask_out[i] = static_cast<float>(y);
        }
    }
    for (size_t k = 0; k < half; k++) {
        mask_out[nfft - 1 - k] = mask_out[k];
    }

    for (size_t k = 0; k < nfft; k++) {
        mask_out[k] = mask_out[k] * (fVar30 * 1.0f) + (1.0f - fVar30);
    }

    // RT_DUMP_BIN: single-frame per-bin tract at frame RT_DUMP_FRAME (default
    // 100): k am res lvl_raw band_level post-IIR1/2, pre-warp mask, W weight.
    if (dbg_path && !dbg_prewarp.empty()) {
        static int dbg_frames = 0;
        int dbg_target = 100;
        if (const char* fs = getenv("RT_DUMP_FRAME")) dbg_target = atoi(fs);
        if (dbg_frames++ != dbg_target) return;
        FILE* df = fopen(dbg_path, "wb");
        if (df) {
            fprintf(df, "# fc=%g q=%g sens=%g rp=%.6f\n", band.fc, band.q, band.sens, rp);
            for (size_t k = 0; k < nbin; k++) {
                double res_k = std::max(static_cast<double>(res[k]), 1e-12);
                double lvl_raw = static_cast<double>(am[k]) / res_k * scale_factor;
                double w = kBand768[k] * kWarp[k] * std::pow(res_k, rp);
                fprintf(df, "%zu %.9g %.9g %.9g %.9g %.9g %.9g\n", k,
                        static_cast<double>(am[k]), res_k, lvl_raw,
                        static_cast<double>(band_level[k]),
                        static_cast<double>(dbg_prewarp[k]), w);
            }
            fclose(df);
        }
    }

    // RT_DUMP_ALL trajectory: append per-frame lvl_raw spectrum (binary:
    // int32 frame, int32 nbin, float32 lvl_raw[nbin]). Single-band cases only.
    // Detector path is law-independent -> one capture serves offline law fits.
    if (const char* ap = getenv("RT_DUMP_ALL")) {
        static FILE* af = fopen(ap, "ab");
        if (af) {
            static int aframe = 0;
            int32_t hdr[2] = {static_cast<int32_t>(aframe++),
                              static_cast<int32_t>(nbin)};
            fwrite(hdr, sizeof(int32_t), 2, af);
            for (size_t k = 0; k < nbin; k++) {
                double res_k = std::max(static_cast<double>(res[k]), 1e-12);
                float lv = static_cast<float>(
                    static_cast<double>(am[k]) / res_k * scale_factor);
                fwrite(&lv, sizeof(float), 1, af);
            }
            fflush(af);
        }
    }
}

}  // 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);
        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, 1.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);
    }

    for (size_t k = 0; k <= half; k++) mask[k] = 1.0f;

    if (is_internal_grid(nfft_, sample_rate_)) {
        for (size_t b = 0; b < bands_.size(); b++) {
            std::vector<float> band_mask(nfft_, 1.0f);
            process_band_structural(am_.data(), res_[b].data(), bands_[b],
                                    band_mask.data(), nfft_, sample_rate_);
            for (size_t k = 0; k <= half; k++) {
                mask[k] = std::min(band_mask[k], mask[k]);
            }
        }
    } else {
        for (size_t b = 0; b < bands_.size(); b++) {
            double rp = RP0 * std::pow(static_cast<double>(bands_[b].q), DRP);
            double fk = 0.0;
            double fstep = (sample_rate_ * 0.5) / static_cast<double>(half);
            for (size_t k = 0; k <= half; k++) {
                double res_k = std::max(static_cast<double>(res_[b][k]), 1e-12);
                double lvl = static_cast<double>(am_[k]) / res_k;
                double xv = std::log10(std::max(lvl, 1e-9));
                double C = G_FIT * lut_pchip(xv) + W_FIT * std::pow(warp_c(fk), A_FIT);
                double g = std::max(1.0 - C, 1e-9) * std::pow(res_k, rp);
                mask[k] = std::min(static_cast<float>(g), mask[k]);
                fk += fstep;
            }
        }
    }

    for (size_t k = half + 1; k < nfft_; k++) {
        mask[k] = mask[nfft_ - k];
    }
}