#include <psemek/audio/effect/filter.hpp>

namespace psemek::audio
{

	namespace
	{

		struct linear_filter_impl
			: stream
		{
			linear_filter_impl(stream_ptr stream, util::array<float> a, util::array<float> b)
				: stream_(std::move(stream))
				, a_(std::move(a))
				, b_(std::move(b))
			{
				prev_in_[0].resize(b_.size(), 0.f);
				prev_in_[1].resize(b_.size(), 0.f);
				prev_out_[0].resize(a_.size(), 0.f);
				prev_out_[1].resize(a_.size(), 0.f);
			}

			std::optional<std::size_t> length() const override
			{
				return stream_->length();
			}

			std::size_t read(util::span<float> samples) override
			{
				std::size_t count = stream_->read(samples);

				// TODO: optimize? Maybe split into 3 loops (first N, main part, last N) to
				// get rid of shift_right on each iteration
				for (std::size_t i = 0; i < count; i += 2)
				{
					for (std::size_t s : {0, 1})
					{
						if (!prev_in_[s].empty())
						{
							std::shift_right(prev_in_[s].begin(), prev_in_[s].end(), 1);
							prev_in_[s][0] = samples[i + s];
						}

						float value = 0.f;
						for (std::size_t j = 0; j < b_.size(); ++j)
							value += prev_in_[s][j] * b_[j];
						for (std::size_t j = 0; j < a_.size(); ++j)
							value -= prev_out_[s][j] * a_[j];

						samples[i + s] = value;

						if (!prev_out_[s].empty())
						{
							std::shift_right(prev_out_[s].begin(), prev_out_[s].end(), 1);
							prev_out_[s][0] = value;
						}
					}
				}

				return count;
			}

			std::size_t played() const override
			{
				return stream_->played();
			}

		private:
			stream_ptr stream_;
			util::array<float> a_;
			util::array<float> b_;
			util::array<float> prev_in_[2];
			util::array<float> prev_out_[2];
		};

	}

	stream_ptr linear_filter(stream_ptr stream, util::array<float> a, util::array<float> b)
	{
		return std::make_shared<linear_filter_impl>(std::move(stream), std::move(a), std::move(b));
	}

	stream_ptr low_pass_filter(stream_ptr stream)
	{
		return linear_filter(std::move(stream), {}, {1.f, 1.f});
	}

	stream_ptr high_pass_filter(stream_ptr stream)
	{
		return linear_filter(std::move(stream), {}, {1.f, -1.f});
	}

}
