#include "fnfaith.hpp"
#include <cmath>
#include <cstdlib>

namespace fnfaith {

Params params_from_env() {
    Params p{};
    p.scale = getenv("RT_FAITH_SCALE") ? atof(getenv("RT_FAITH_SCALE")) : 1.0;
    p.kappa = getenv("RT_FAITH_KAPPA") ? atof(getenv("RT_FAITH_KAPPA")) : 0.0094;
    p.p     = getenv("RT_FAITH_P")     ? atof(getenv("RT_FAITH_P"))     : 0.822;
    p.tau1  = getenv("RT_FAITH_TAU1")  ? atof(getenv("RT_FAITH_TAU1"))  : 19.03;
    p.mult1 = getenv("RT_FAITH_MULT1") ? atof(getenv("RT_FAITH_MULT1")) : 360.0;
    p.tau3  = getenv("RT_FAITH_TAU3")  ? atof(getenv("RT_FAITH_TAU3"))  : 10.68;
    p.mult3 = getenv("RT_FAITH_MULT3") ? atof(getenv("RT_FAITH_MULT3")) : 2400.0;
    p.tau4  = getenv("RT_FAITH_TAU4")  ? atof(getenv("RT_FAITH_TAU4"))  : 10.68;
    p.mult4 = getenv("RT_FAITH_MULT4") ? atof(getenv("RT_FAITH_MULT4")) : 2400.0;
    p.C_hz  = getenv("RT_FAITH_C")     ? atof(getenv("RT_FAITH_C"))     : 1000.0;
    return p;
}

namespace {

struct FaithCoefs {
    std::vector<double> up, down;
};

// FUN_180533340 transcription (constants from dump: DAT_24c3d8c=0.5, DAT_24c46b8=-2pi)
FaithCoefs gen_coefs(size_t nbin, double sr, double tau, double mult, double p,
                     double C_hz, double kappa) {
    FaithCoefs c;
    c.up.assign(nbin, 0.0);
    c.down.assign(nbin, 1.0);
    const double srh = sr * 0.5;
    const double fcnorm = (C_hz / srh) * (double)nbin;
    for (size_t i = 1; i < nbin; i++) {
        const double r = fcnorm / (double)i;
        const double g = ((double)i <= fcnorm) ? r : std::pow(r, p);
        const double cd = 1.0 / (g * tau / mult + 1.0);
        double up = std::exp(cd * g * tau * kappa * (-6.283185307179586));
        if (up > 1.0) up = 1.0;
        if (up < 0.0) up = 0.0;
        c.up[i] = up;
        c.down[i] = 1.0 - up;
    }
    return c;
}

// FUN_18052d650: acc=0; forward all bins; backward n-2..1, acc persists into bwd.
void bidir(std::vector<double>& accst, const FaithCoefs& cf, float* x, size_t nbin) {
    double acc = 0.0;
    accst[0] = 0.0;
    for (size_t i = 0; i < nbin; i++) {
        acc = cf.up[i] * acc + cf.down[i] * (double)x[i];
        x[i] = (float)acc;
    }
    for (size_t i = (size_t)((long long)nbin - 2); i >= 1; i--) {
        acc = cf.up[i] * acc + cf.down[i] * (double)x[i];
        x[i] = (float)acc;
    }
}

}  // namespace

void band_mask_faithful(const float* am, const float* res, size_t nbin,
                        float sample_rate, float scale_factor, const Params& pr,
                        float* mask_out) {
    static std::vector<FaithCoefs> cache;
    static size_t cached_nbin = 0;
    static double cached_sr = 0.0;
    static Params cached_pr{};
    if (cache.empty() || cached_nbin != nbin || cached_sr != sample_rate ||
        cached_pr.kappa != pr.kappa || cached_pr.p != pr.p ||
        cached_pr.tau1 != pr.tau1 || cached_pr.mult1 != pr.mult1 ||
        cached_pr.tau3 != pr.tau3 || cached_pr.mult3 != pr.mult3 ||
        cached_pr.tau4 != pr.tau4 || cached_pr.mult4 != pr.mult4 ||
        cached_pr.C_hz != pr.C_hz) {
        cache.clear();
        cache.push_back(gen_coefs(nbin, sample_rate, pr.tau1, pr.mult1, pr.p, pr.C_hz, pr.kappa));
        cache.push_back(gen_coefs(nbin, sample_rate, pr.tau3, pr.mult3, pr.p, pr.C_hz, pr.kappa));
        cache.push_back(gen_coefs(nbin, sample_rate, pr.tau4, pr.mult4, pr.p, pr.C_hz, pr.kappa));
        cached_nbin = nbin;
        cached_sr = sample_rate;
        cached_pr = pr;
    }

    // band curve: lvl_raw scaled
    std::vector<float> m(nbin);
    for (size_t k = 0; k < nbin; k++) {
        double res_k = res[k] > 1e-12 ? (double)res[k] : 1e-12;
        double lvl = (double)am[k] / res_k * (double)scale_factor;
        m[k] = (float)(std::exp2(-(lvl * pr.scale)) );
    }

    std::vector<double> accst(1, 0.0);
    bidir(accst, cache[0], m.data(), nbin);   // #1
    bidir(accst, cache[0], m.data(), nbin);   // #2
    bidir(accst, cache[1], m.data(), nbin);   // #3
    bidir(accst, cache[2], m.data(), nbin);   // #4a
    bidir(accst, cache[2], m.data(), nbin);   // #4b

    for (size_t k = 0; k < nbin; k++) mask_out[k] = m[k];
}

}  // namespace fnfaith