#include "AudioBuffer.h"
#include <immintrin.h>
#include "MaxElement.h"
#include "Mixer/MixerTrack.h"

void FAudioBuffer::Resize(int32 NumChannels, int32 NumSamples)
{
	UE::TScopeLock Lock(CriticalSection);
	Buffer.SetNum(NumChannels);
	Headers.SetNum(NumChannels);
	Peak.SetNum(NumChannels);
	for (int32 i = 0; i < NumChannels; i++)
	{
		Buffer[i].SetNum(NumSamples);
		Headers[i] = Buffer[i].GetData();
	}
}

void FAudioBuffer::ZeroBuffer()
{
	UE::TScopeLock Lock(CriticalSection);
	for (int32 i = 0; i < Buffer.Num(); i++)
	{
		FMemory::Memzero(Buffer[i].GetData(), Buffer[i].Num() * sizeof(float));
	}
}

void FAudioBuffer::AddBuffer(FAudioBuffer& Other, float Gain)
{
	UE::TScopeLock Lock(CriticalSection);
	// AVX 256-bit
	for (int32 Channel = 0; Channel < Buffer.Num(); ++Channel)
	{
		TArray<float>& ChannelBuffer = Buffer[Channel];
		for (int32 Sample = 0; Sample < ChannelBuffer.Num(); Sample += 8)
		{
			const __m256 a = _mm256_load_ps(ChannelBuffer.GetData() + Sample);
			const __m256 b = _mm256_load_ps(Other.Buffer[Channel].GetData() + Sample);
			const __m256 c = _mm256_add_ps(a, b);
			const __m256 d = _mm256_mul_ps(c, _mm256_set1_ps(Gain));
			_mm256_store_ps(ChannelBuffer.GetData() + Sample, d);
		}
	}
}

void FAudioBuffer::CalculatePeak()
{
	UE::TScopeLock Lock(CriticalSection);
	for (int i = 0; i < Buffer.Num(); ++i)
	{
		Peak[i] = *Algo::MaxElement(Buffer[i]);
	}
}

float FAudioBuffer::GetPeak(int32 Channel)
{
	return Peak[Channel];
}

void FAudioBuffer::Multiple(float In)
{
	UE::TScopeLock Lock(CriticalSection);
	for (int32 Channel = 0; Channel < Buffer.Num(); ++Channel)
	{
		TArray<float>& ChannelBuffer = Buffer[Channel];
		for (int32 Sample = 0; Sample < ChannelBuffer.Num(); Sample += 8)
		{
			const __m256 a = _mm256_load_ps(ChannelBuffer.GetData() + Sample);
			const __m256 b = _mm256_set1_ps(In);
			const __m256 c = _mm256_mul_ps(a, b);
			_mm256_store_ps(ChannelBuffer.GetData() + Sample, c);
		}
	}
}