#include "filter.hpp"
#include <cmath>
#include <cstring>

DigitalFilter::DigitalFilter() {
    memset(&state_, 0, sizeof(state_));
    b0 = 1; b1 = 0; b2 = 0;
    a1 = 0; a2 = 0;
}

void DigitalFilter::setParams(const FilterParams& p, float sample_rate) {
    params_ = p;
    fs_ = sample_rate;
    updateCoeffs();
}

void DigitalFilter::updateCoeffs() {
    if (!params_.on) {
        b0 = 1; b1 = 0; b2 = 0;
        a1 = 0; a2 = 0;
        return;
    }
    
    float wc = 2 * M_PI * params_.freq / fs_;
    float tan_wc = std::tan(wc / 2);
    float cos_wc = std::cos(wc);
    
    if (params_.type == 0) {  // peak
        float Q = params_.q;
        float alpha = tan_wc / (2 * Q);
        float k = std::pow(10, params_.gain / 40);
        
        b0 = 1 + alpha * k;
        b1 = -2 * cos_wc;
        b2 = 1 - alpha * k;
        float a0_inv = 1 / (1 + alpha);
        b0 *= a0_inv; b1 *= a0_inv; b2 *= a0_inv;
        a1 = -2 * cos_wc * a0_inv;
        a2 = -(1 - alpha) * a0_inv;
    } else if (params_.type == 1) {  // shelf
        float Q = params_.q;
        float A = std::pow(10, params_.gain / 40);
        float alpha = tan_wc / (2 * Q);
        float beta = std::sqrt(A);
        
        if (params_.gain >= 0) {
            b0 = A * ((A + 1) + (A - 1) * cos_wc + 2 * beta * tan_wc);
            b2 = A * ((A + 1) + (A - 1) * cos_wc - 2 * beta * tan_wc);
        } else {
            b0 = (A + 1) - (A - 1) * cos_wc + 2 * beta * tan_wc;
            b2 = (A + 1) - (A - 1) * cos_wc - 2 * beta * tan_wc;
        }
        float a0_inv = 1 / ((A + 1) - (A - 1) * cos_wc + 2 * beta * tan_wc);
        b0 *= a0_inv; b2 *= a0_inv;
        a1 = -2 * ((A - 1) - (A + 1) * cos_wc) * a0_inv;
        a2 = -((A + 1) - (A - 1) * cos_wc - 2 * beta * tan_wc) * a0_inv;
    } else if (params_.type == 2) {  // reject
        float Q = params_.q;
        float alpha = tan_wc / (2 * Q);
        
        b0 = 1;
        b1 = -2 * cos_wc;
        b2 = 1;
        float a0_inv = 1 / (1 + alpha);
        b1 *= a0_inv; b2 *= a0_inv;
        a1 = -2 * cos_wc * a0_inv;
        a2 = -(1 - alpha) * a0_inv;
    }
}

float DigitalFilter::process(float x) {
    float y = b0 * x + b1 * state_.x1 + b2 * state_.x2 - a1 * state_.y1 - a2 * state_.y2;
    state_.x2 = state_.x1;
    state_.x1 = x;
    state_.y2 = state_.y1;
    state_.y1 = y;
    return y;
}

FilterGraph::FilterGraph() {
}

void FilterGraph::processBlock(float* in, float* out, size_t n, const FilterParams* bands, size_t num_bands) {
    for (size_t i = 0; i < n; i++) {
        out[i] = in[i];
    }
    for (size_t b = 0; b < num_bands; b++) {
        if (!bands[b].on) continue;
        for (size_t i = 0; i < n; i++) {
            out[i] = filters_[b].process(out[i]);
        }
    }
}