/******

Copyright (C) 2024 Polyspectral LLC

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
******/

// This code implements a technique for optimized IIR filters, which is described
//  in an article at https://link.polyspectral.com/articles-efficientiirs

#include <iostream>
#include <random>
#include <chrono>

#ifdef _MSC_VER
# define RESTRICT __restrict
#else
# define RESTRICT
#endif

// straightforward transposed direct form II implementation
//  (see https://ccrma.stanford.edu/~jos/fp/Transposed_Direct_Forms.html)

struct SimpleBiquad
{
   SimpleBiquad( float b0, float b1, float b2,
                 float a0, float a1, float a2 )
      : m_b0(b0 / a0)
      , m_b1(b1 / a0)
      , m_b2(b2 / a0)
      , m_a1(a1 / a0)
      , m_a2(a2 / a0)
   {
   }

   void Process(const float* input, float* output, size_t numSamples)
   {
      // move these into local variables for efficiency
      float b0 = m_b0;
      float b1 = m_b1;
      float b2 = m_b2;
      float a1 = m_a1;
      float a2 = m_a2;
      float s1 = m_s1;
      float s2 = m_s2;

      for (size_t i = 0; i < numSamples; ++i)
      {
         float x = input[i];
         float y = b0 * x + s1;
         s1 = b1 * x - a1 * y + s2;
         s2 = b2 * x - a2 * y;
         output[i] = y;
      }

      m_s1 = s1;
      m_s2 = s2;
   }

   float m_b0, m_b1, m_b2, m_a1, m_a2;
   float m_s1{}, m_s2{};
};

// matrix biquad implementation

template<int N>
struct MatrixBiquad
{
   MatrixBiquad( float b0, float b1, float b2,
                 float a0, float a1, float a2 )
      : m_b0(b0 / a0)
      , m_b1(b1 / a0)
      , m_b2(b2 / a0)
      , m_a1(a1 / a0)
      , m_a2(a2 / a0)
   {
      float input[N]{};
      float output[N]{};

      SimpleBiquad bq(b0, b1, b2, a0, a1, a2);

      // record filter impulse responses into the first N columns of
      //  the matrix. note that this can be done more efficiently by
      //  processing a single impulse step-by-step -- we do it this
      //  more inefficient way for the sake of clarity.
      for (int i = 0; i < N; ++i)
      {
         // set up an impulse input at sample i
         bq.m_s1 = bq.m_s2 = 0;
         for (int j = 0; j < N; ++j)
            input[j] = (i == j ? 1.0f : 0.0f);

         bq.Process(input, output, N);

         // record the results into column i
         for (int j = 0; j < N; ++j)
            m_matrix[i][j] = output[j];
         m_matrix[i][N] = bq.m_s1;
         m_matrix[i][N + 1] = bq.m_s2;
      }

      // zero out the input
      for (int j = 0; j < N; ++j)
          input[j] = 0.0f;

      // record the response to s1=1 into column N
      bq.m_s1 = 1;
      bq.m_s2 = 0;
      bq.Process(input, output, N);
      for (int i = 0; i < N; ++i)
         m_matrix[N][i] = output[i];
      m_matrix[N][N] = bq.m_s1;
      m_matrix[N][N + 1] = bq.m_s2;

      // record the response to s2=1 into column N+1
      bq.m_s1 = 0;
      bq.m_s2 = 1;
      bq.Process(input, output, N);
      for (int i = 0; i < N; ++i)
         m_matrix[N + 1][i] = output[i];
      m_matrix[N + 1][N] = bq.m_s1;
      m_matrix[N + 1][N + 1] = bq.m_s2;
   }

   void Process(const float* RESTRICT input, float* RESTRICT output, size_t numSamples)
   {
      float s1 = m_s1;
      float s2 = m_s2;
      float y[N + 2];

      size_t samp = 0;
      while (samp + N <= numSamples)
      {
         // clear output vector
         for (int j = 0; j < N + 2; ++j)
            y[j] = 0;

         // multiply the first N columns by consecutive input samples,
         //  add to the output vector
         float* RESTRICT column;
         for (int i = 0; i < N; ++i)
         {
            float x = input[samp++];
            column = &m_matrix[i][0];
            for (int j = 0; j < N + 2; ++j)
               y[j] += column[j] * x;
         }

         // multiply the final two columns by s1 and s2
         column = &m_matrix[N][0];
         for (int j = 0; j < N + 2; ++j)
            y[j] += column[j] * s1;
         column = &m_matrix[N + 1][0]; 
         for (int j = 0; j < N + 2; ++j)
            y[j] += column[j] * s2;

         // store output samples
         for (int j = 0; j < N; ++j)
            output[j] = y[j];

         // read the updated state variables out of the output vector
         s1 = y[N];
         s2 = y[N + 1];

         output += N;
      }

      // process any leftover samples the usual way
      for (; samp < numSamples; ++samp, ++output)
      {
         float x = input[samp];
         float y = m_b0 * x + s1;
         s1 = m_b1 * x - m_a1 * y + s2;
         s2 = m_b2 * x - m_a2 * y;
         *output = y;
      }

      m_s1 = s1;
      m_s2 = s2;
   }

   float m_matrix[N + 2][N + 2] {};

   float m_b0, m_b1, m_b2, m_a1, m_a2;
   float m_s1 = 0;
   float m_s2 = 0;
};

int main()
{
   const size_t kSize = 1000000;
   std::vector<float> input(kSize), output1(kSize), output2(kSize);

   // fill input with white noise
   {
      std::random_device rd;
      std::mt19937 gen(rd());
      std::uniform_real_distribution<float> rnd(-1, 1);
      for (auto itr = input.begin(); itr != input.end(); ++itr)
         *itr = rnd(gen);
   }

   // example low-pass coefficients
   const float
      b0 = 0.05f,
      b1 = 0.1f,
      b2 = 0.05f,
      a0 = 1.5f,
      a1 = -01.8f,
      a2 = 0.5f;

   // run through simple biquad
   auto t1 = std::chrono::high_resolution_clock::now();
   SimpleBiquad bq1(b0, b1, b2, a0, a1, a2);
   bq1.Process(input.data(), output1.data(), input.size());
   auto t2 = std::chrono::high_resolution_clock::now();
   std::cout << "SimpleBiquad: " << std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() << "us\n";

   // run through matrix biquad
   t1 = std::chrono::high_resolution_clock::now();
   MatrixBiquad<6> bq6(b0, b1, b2, a0, a1, a2);
   bq6.Process(input.data(), output2.data(), input.size());
   t2 = std::chrono::high_resolution_clock::now();
   std::cout << "MatrixBiquad<6>: " << std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() << "us\n";

   // check for approximate equality
   float largestError = 0;
   for (size_t i = 0; i < kSize; ++i)
      largestError = std::max(largestError, std::abs(output1[i] - output2[i]));
   std::cout << "\nError magnitude: " << 20 * std::log10(largestError) << "dB\n";

   return 0;
}