#include "HyperTwistRecognition/HyperTwistRecognitionReplayLibrary.h"

#include "JsonObjectConverter.h"

namespace HyperTwistRecognitionReplayLibraryInternal
{
	template <typename TStruct>
	bool DeserializePayload(const FString& Json, TStruct& OutValue)
	{
		if (Json.IsEmpty())
		{
			return false;
		}

		return FJsonObjectConverter::JsonObjectStringToUStruct(Json, &OutValue, 0, 0);
	}

	FString ResolveRecognitionStage(const FHyperTwistReplayRecognitionPayload& Payload)
	{
		if (!Payload.ScanStage.IsEmpty())
		{
			return Payload.ScanStage;
		}

		return Payload.CommittedFaceOrStage;
	}

	void StartEpisodeFromEvent(
		FHyperTwistRecognitionReplayEpisode& Episode,
		const FHyperTwistReplayEvent& Event,
		const FHyperTwistReplayRecognitionPayload& Payload,
		int32 EpisodeOrdinal)
	{
		Episode = FHyperTwistRecognitionReplayEpisode();
		Episode.EpisodeId = FString::Printf(TEXT("rec_ep_%03d"), EpisodeOrdinal);
		Episode.Stage = ResolveRecognitionStage(Payload);
		Episode.StartedAtMs = Event.TimeMs;
	}

	void FoldConfidence(FHyperTwistRecognitionReplayEpisode& Episode, const float Confidence, int32& InOutConfidenceCount, float& InOutConfidenceSum)
	{
		InOutConfidenceSum += Confidence;
		++InOutConfidenceCount;

		if (Episode.PeakConfidence < 0.0f)
		{
			Episode.PeakConfidence = Confidence;
		}
		else
		{
			Episode.PeakConfidence = FMath::Max(Episode.PeakConfidence, Confidence);
		}
	}

	void FinalizeEpisode(
		FHyperTwistRecognitionReplayEpisode& Episode,
		const int32 ClosedAtMs,
		const bool bCompleted,
		TArray<FHyperTwistRecognitionReplayEpisode>& OutEpisodes)
	{
		Episode.CompletedAtMs = ClosedAtMs;
		Episode.DurationMs = FMath::Max(0, Episode.CompletedAtMs - Episode.StartedAtMs);
		Episode.bCompleted = bCompleted;
		Episode.bAmbiguous = Episode.CorrectionCount > 0 || !Episode.bCompleted;
		OutEpisodes.Add(Episode);
		Episode = FHyperTwistRecognitionReplayEpisode();
	}
}

TArray<FHyperTwistRecognitionReplayEpisode> UHyperTwistRecognitionReplayLibrary::CollectRecognitionReplayEpisodes(const FHyperTwistReplayPacket& Packet)
{
	TArray<FHyperTwistRecognitionReplayEpisode> Episodes;
	FHyperTwistRecognitionReplayEpisode OpenEpisode;
	bool bHasOpenEpisode = false;
	int32 EpisodeOrdinal = 1;
	int32 ConfidenceCount = 0;
	float ConfidenceSum = 0.0f;

	for (const FHyperTwistReplayEvent& Event : Packet.Events)
	{
		switch (Event.EventType)
		{
		case EHyperTwistReplayEventType::RecognitionPreview:
		case EHyperTwistReplayEventType::RecognitionScan:
		case EHyperTwistReplayEventType::RecognitionCommit:
		case EHyperTwistReplayEventType::RecognitionCorrection:
		{
			FHyperTwistReplayRecognitionPayload Payload;
			if (!HyperTwistRecognitionReplayLibraryInternal::DeserializePayload(Event.PayloadJson, Payload))
			{
				break;
			}

			const FString EventStage = HyperTwistRecognitionReplayLibraryInternal::ResolveRecognitionStage(Payload);
			if (!bHasOpenEpisode)
			{
				HyperTwistRecognitionReplayLibraryInternal::StartEpisodeFromEvent(OpenEpisode, Event, Payload, EpisodeOrdinal++);
				bHasOpenEpisode = true;
				ConfidenceCount = 0;
				ConfidenceSum = 0.0f;
			}
			else if (!EventStage.IsEmpty() && !OpenEpisode.Stage.IsEmpty() && OpenEpisode.Stage != EventStage)
			{
				OpenEpisode.AverageConfidence = ConfidenceCount > 0 ? ConfidenceSum / static_cast<float>(ConfidenceCount) : -1.0f;
				HyperTwistRecognitionReplayLibraryInternal::FinalizeEpisode(OpenEpisode, Event.TimeMs, false, Episodes);
				HyperTwistRecognitionReplayLibraryInternal::StartEpisodeFromEvent(OpenEpisode, Event, Payload, EpisodeOrdinal++);
				ConfidenceCount = 0;
				ConfidenceSum = 0.0f;
			}
			else if (OpenEpisode.Stage.IsEmpty())
			{
				OpenEpisode.Stage = EventStage;
			}

			HyperTwistRecognitionReplayLibraryInternal::FoldConfidence(OpenEpisode, Payload.Confidence, ConfidenceCount, ConfidenceSum);

			if (Event.EventType == EHyperTwistReplayEventType::RecognitionPreview)
			{
				++OpenEpisode.PreviewCount;
			}
			else if (Event.EventType == EHyperTwistReplayEventType::RecognitionScan)
			{
				++OpenEpisode.ScanCount;
			}
			else if (Event.EventType == EHyperTwistReplayEventType::RecognitionCorrection)
			{
				++OpenEpisode.CorrectionCount;
			}
			else
			{
				OpenEpisode.FinalCommitConfidence = Payload.Confidence;
				OpenEpisode.FinalCommittedUnit = Payload.CommittedFaceOrStage;
				OpenEpisode.AverageConfidence = ConfidenceCount > 0 ? ConfidenceSum / static_cast<float>(ConfidenceCount) : -1.0f;
				HyperTwistRecognitionReplayLibraryInternal::FinalizeEpisode(OpenEpisode, Event.TimeMs, true, Episodes);
				bHasOpenEpisode = false;
				ConfidenceCount = 0;
				ConfidenceSum = 0.0f;
			}
			break;
		}

		default:
			break;
		}
	}

	if (bHasOpenEpisode)
	{
		const int32 ClosedAtMs = Packet.Events.Num() > 0 ? Packet.Events.Last().TimeMs : OpenEpisode.StartedAtMs;
		OpenEpisode.AverageConfidence = ConfidenceCount > 0 ? ConfidenceSum / static_cast<float>(ConfidenceCount) : -1.0f;
		HyperTwistRecognitionReplayLibraryInternal::FinalizeEpisode(OpenEpisode, ClosedAtMs, false, Episodes);
	}

	return Episodes;
}

FHyperTwistRecognitionReplaySummary UHyperTwistRecognitionReplayLibrary::DeriveRecognitionReplaySummary(const FHyperTwistReplayPacket& Packet)
{
	FHyperTwistRecognitionReplaySummary Summary;
	Summary.Episodes = CollectRecognitionReplayEpisodes(Packet);
	Summary.EpisodeCount = Summary.Episodes.Num();

	TSet<FString> CoveredStages;
	int32 DurationSum = 0;
	int32 CompletedDurationCount = 0;
	float CommitConfidenceSum = 0.0f;
	int32 CommitConfidenceCount = 0;

	for (const FHyperTwistReplayEvent& Event : Packet.Events)
	{
		switch (Event.EventType)
		{
		case EHyperTwistReplayEventType::RecognitionPreview:
			++Summary.TotalRecognitionEventCount;
			++Summary.TotalPreviewCount;
			break;

		case EHyperTwistReplayEventType::RecognitionScan:
			++Summary.TotalRecognitionEventCount;
			++Summary.TotalScanCount;
			break;

		case EHyperTwistReplayEventType::RecognitionCommit:
			++Summary.TotalRecognitionEventCount;
			++Summary.TotalCommitCount;
			break;

		case EHyperTwistReplayEventType::RecognitionCorrection:
			++Summary.TotalRecognitionEventCount;
			++Summary.TotalCorrectionCount;
			break;

		default:
			break;
		}
	}

	for (const FHyperTwistRecognitionReplayEpisode& Episode : Summary.Episodes)
	{
		if (!Episode.Stage.IsEmpty())
		{
			CoveredStages.Add(Episode.Stage);
		}

		if (Episode.bCompleted)
		{
			++Summary.CompletedEpisodeCount;
			DurationSum += Episode.DurationMs;
			++CompletedDurationCount;

			if (Episode.FinalCommitConfidence >= 0.0f)
			{
				CommitConfidenceSum += Episode.FinalCommitConfidence;
				++CommitConfidenceCount;
			}

			Summary.LastCommittedUnit = Episode.FinalCommittedUnit;
		}

		if (Episode.bAmbiguous)
		{
			++Summary.AmbiguousEpisodeCount;
		}
	}

	for (const FString& CoveredStage : CoveredStages)
	{
		Summary.CoveredStages.Add(CoveredStage);
	}
	Summary.CoveredStages.Sort();
	Summary.AverageEpisodeDurationMs = CompletedDurationCount > 0 ? DurationSum / CompletedDurationCount : 0;
	Summary.AverageCommitConfidence = CommitConfidenceCount > 0 ? CommitConfidenceSum / static_cast<float>(CommitConfidenceCount) : -1.0f;

	return Summary;
}
