#include "WaveformViewer.h"

#include "SlatePixelShader.h"

SWaveformViewer::SWaveformViewer()
{
	Param.WaveColor = FLinearColor(0.0f, 1.0f, 0.0f, 1.0f).ToFColorSRGB();
	Param.BgColor = FLinearColor(0.0f, 0.0f, 0.0f, 0.0f).ToFColorSRGB();
	Param.LineUV = 0.0f;
}

void SWaveformViewer::Construct(const FArguments& InArgs)
{
	WaveformHandle = InArgs._WaveformHandle;
	
    FSlateRenderer* Renderer = FSlateApplication::Get().GetRenderer();
	
	RenderTarget = Renderer->CreateRenderTargetTexture(1, 1, EPixelFormat::PF_R8G8B8A8);
	
	TArray<float> InitData;
	InitData.SetNumZeroed(2);
	if (Renderer->GetRendererAPI() == ERendererAPI::D3D11)
	{
		Params = Renderer->CreateComputeShaderParam(ComputeShader);
		Params->Init();
		Params->AddComputeShaderTextureParameter("Result", RenderTarget, 2);
		Params->AddComputeShaderParam("Samples", InitData.GetData(), InitData.Num(), 0);
		Params->AddShaderParam("Params", (bool*)&Param, sizeof(FCSParam) / sizeof(bool));

		ComputeShader = Renderer->CreateComputeShader(TEXT("Arona"), TEXT("WaveformCS"));
	}
	else if (Renderer->GetRendererAPI() == ERendererAPI::OpenGL)
	{
		ComputeShader = Renderer->CreateComputeShader(TEXT("Arona"), TEXT("WaveformCS"));

		Params = Renderer->CreateComputeShaderParam(ComputeShader);
		Params->Init();
		Params->AddComputeShaderTextureParameter("Result", RenderTarget, 2);
		Params->AddComputeShaderParam("Samples", InitData.GetData(), InitData.Num(), 0);
		Params->AddComputeShaderParam("Params", (bool*)&Param, sizeof(FCSParam) / sizeof(bool), 1);
	}
	
 	ComputeElement = Renderer->MakeComputeShaderElement(ComputeShader, Params);
}

void SWaveformViewer::Tick(const FGeometry& AllottedGeometry, const double InCurrentTime, const float InDeltaTime)
{
	SLeafWidget::Tick(AllottedGeometry, InCurrentTime, InDeltaTime);
}

TArray<FSamplePeak> ScaleWaveform(const TArrayView<const FSamplePeak>& Waveform, const int32 SizeX)
{
	TArray<FSamplePeak> Out;
	Out.SetNum(SizeX);
	const float Scale = (float)Waveform.Num() / SizeX;
	float Index = 0;
	for (int32 i = 0; i < SizeX; ++i)
	{
		const int32 Begin = FMath::FloorToInt(Index);
		const int32 End = FMath::Min(FMath::CeilToInt(Index + Scale), Waveform.Num() - 1);
		FSamplePeak& Peak = Out[i];
		for (int32 j = Begin; j < End; ++j)
		{
			Peak += Waveform[j];
		}
		Index += Scale;
	}
	return Out;
}

int32 SWaveformViewer::OnPaint(const FPaintArgs& Args, const FGeometry& AllottedGeometry,
                               const FSlateRect& MyCullingRect, FSlateWindowElementList& OutDrawElements, int32 LayerId,
                               const FWidgetStyle& InWidgetStyle, bool bParentEnabled) const
{
	if (WaveformHandle.IsValid())
	{
		const FIntPoint& Size = FIntPoint(AllottedGeometry.Size.X, AllottedGeometry.Size.Y);
		if (Size.X <= 0 || Size.Y <= 0) return LayerId;
		const TArrayView<const FSamplePeak>& SamplePeaks = WaveformHandle->GetWaveform(Size.X);
		if (!SamplePeaks.Num()) return LayerId;
		Param.LineUV = 1.0f / Size.Y;
		Param.PeaksPerPixel = (float)SamplePeaks.Num() / Size.X;
		const TArray<FSamplePeak>& Peaks = ScaleWaveform(SamplePeaks, Size.X);
		Params->AddComputeShaderParam("Samples", (float*)Peaks.GetData(), Peaks.Num() * 2, 0);
		
		if (FSlateApplication::Get().GetRenderer()->GetRendererAPI() == ERendererAPI::OpenGL)
			Params->AddComputeShaderParam("Params", (bool*)&Param, sizeof(FCSParam) / sizeof(bool), 1);
		else
			Params->AddShaderParam("Params", (bool*)&Param, sizeof(FCSParam) / sizeof(bool));
		
		RenderTarget->Resize(Size);
		FSlateDrawElement::MakeComputeShaderDispatch(OutDrawElements, LayerId, Size.X , Size.Y, 1, ComputeElement);
	}
	
 	LayerId++;
	FSlateDrawElement::MakeViewport(OutDrawElements, LayerId, AllottedGeometry.ToPaintGeometry(), RenderTarget, ESlateDrawEffect::None, FLinearColor::White);
	return LayerId;
}