#include "VoxelForgeRegionEvalExperiment.h"

#include "Async/ParallelFor.h"
#include "HAL/PlatformTime.h"
#include "VoxelCaveMorphology.h"
#include "VoxelDensityProfile.h"
#include "VoxelGenerator.h"
#include "HAL/PlatformTLS.h"

#if PLATFORM_WINDOWS
#include "Windows/AllowWindowsPlatformTypes.h"
THIRD_PARTY_INCLUDES_START
#include <Windows.h>
THIRD_PARTY_INCLUDES_END
#include "Windows/HideWindowsPlatformTypes.h"
#endif

namespace VoxelForgeRegionEval
{
    namespace
    {
        constexpr double WindowsFileTimeSeconds = 1.0e-7;

        uint64 Fnv1aByte(uint64 Hash, uint8 Byte)
        {
            Hash ^= Byte;
            return Hash * 1099511628211ull;
        }

        uint64 HashRawFloats(const TArray<float>& Densities)
        {
            uint64 Hash = 14695981039346656037ull;
            for (const float Density : Densities)
            {
                uint32 Bits = 0;
                FMemory::Memcpy(&Bits, &Density, sizeof(Bits));
                Hash = Fnv1aByte(Hash, static_cast<uint8>(Bits & 0xFFu));
                Hash = Fnv1aByte(Hash, static_cast<uint8>((Bits >> 8) & 0xFFu));
                Hash = Fnv1aByte(Hash, static_cast<uint8>((Bits >> 16) & 0xFFu));
                Hash = Fnv1aByte(Hash, static_cast<uint8>((Bits >> 24) & 0xFFu));
            }
            return Hash;
        }

        bool GetCurrentThreadCpuSeconds(double& OutSeconds)
        {
#if PLATFORM_WINDOWS
            FILETIME CreationTime = {};
            FILETIME ExitTime = {};
            FILETIME KernelTime = {};
            FILETIME UserTime = {};
            if (GetThreadTimes(
                    GetCurrentThread(), &CreationTime, &ExitTime, &KernelTime, &UserTime) == 0)
            {
                return false;
            }
            ULARGE_INTEGER KernelTicks = {};
            ULARGE_INTEGER UserTicks = {};
            KernelTicks.LowPart = KernelTime.dwLowDateTime;
            KernelTicks.HighPart = KernelTime.dwHighDateTime;
            UserTicks.LowPart = UserTime.dwLowDateTime;
            UserTicks.HighPart = UserTime.dwHighDateTime;
            OutSeconds = static_cast<double>(KernelTicks.QuadPart + UserTicks.QuadPart)
                * WindowsFileTimeSeconds;
            return true;
#else
            (void)OutSeconds;
            return false;
#endif
        }

        int64 LatticeCount(int32 Dimension)
        {
            return static_cast<int64>(Dimension) * Dimension * Dimension;
        }
    }

    int64 LinearIndex(int32 X, int32 Y, int32 Z, int32 Dimension)
    {
        return (static_cast<int64>(Z) * Dimension + Y) * Dimension + X;
    }

    int16 QuantizeDensity(float Density, bool& bOutSaturated, bool& bOutFinite)
    {
        bOutFinite = FMath::IsFinite(Density);
        bOutSaturated = false;
        if (!bOutFinite)
        {
            return 0;
        }
        const double Scaled = static_cast<double>(Density) * QuantizationScale;
        const double Rounded = Scaled >= 0.0
            ? FMath::FloorToDouble(Scaled + 0.5)
            : FMath::CeilToDouble(Scaled - 0.5);
        if (Rounded > static_cast<double>(MAX_int16))
        {
            bOutSaturated = true;
            return MAX_int16;
        }
        if (Rounded < static_cast<double>(MIN_int16))
        {
            bOutSaturated = true;
            return MIN_int16;
        }
        return static_cast<int16>(Rounded);
    }

    uint64 HashQuantized(const TArray<float>& Densities)
    {
        uint64 Hash = 14695981039346656037ull;
        for (const float Density : Densities)
        {
            bool bSaturated = false;
            bool bFinite = false;
            const uint16 Q = static_cast<uint16>(QuantizeDensity(Density, bSaturated, bFinite));
            Hash = Fnv1aByte(Hash, static_cast<uint8>(Q & 0xFFu));
            Hash = Fnv1aByte(Hash, static_cast<uint8>(Q >> 8));
        }
        return Hash;
    }

    bool Evaluate(
        const UVoxelGenerator& Generator,
        FIntVector OriginVoxels,
        int32 Dimension,
        const TArray<int64>* SampleIndices,
        bool bFullRing,
        bool bPrepared,
        float ReachScale,
        TArray<float>* OutDensities,
        FTiming& OutTiming)
    {
        OutTiming = FTiming();
        if (Dimension <= 0 || Dimension > 4096
            || (bFullRing && SampleIndices != nullptr)
            || (!bFullRing && SampleIndices == nullptr))
        {
            return false;
        }
        const int64 TotalLatticeSamples = LatticeCount(Dimension);
        const int64 SampleCount = bFullRing
            ? TotalLatticeSamples : static_cast<int64>(SampleIndices->Num());
        if (SampleCount <= 0 || SampleCount > TotalLatticeSamples)
        {
            return false;
        }
        const double StartSeconds = FPlatformTime::Seconds();
        if (OutDensities != nullptr)
        {
            if (SampleCount > MAX_int32)
            {
                return false;
            }
            OutDensities->SetNumUninitialized(static_cast<int32>(SampleCount));
        }

        const int32 OriginX = OriginVoxels.X;
        const int32 OriginY = OriginVoxels.Y;
        const int32 OriginZ = OriginVoxels.Z;
        double WorkerSeconds[WorkerCount] = {};
        double WorkerCpuSeconds[WorkerCount] = {};
        bool WorkerCpuFromOsThreadTimes[WorkerCount] = {};
        bool bAllWorkerCpuFromOsThreadTimes = true;
        uint32 WorkerThreadIds[WorkerCount] = {};
        uint64 WorkerNonFinite[WorkerCount] = {};
        ParallelFor(WorkerCount, [&](int32 WorkerIndex)
        {
            VoxelCaveMorphology::SetResearchBrickBinningEnabled(bPrepared);
            VoxelCaveMorphology::SetResearchBrickRoomReachScale(ReachScale);
            WorkerThreadIds[WorkerIndex] = FPlatformTLS::GetCurrentThreadId();
            const int64 Begin = (SampleCount * WorkerIndex) / WorkerCount;
            const int64 End = (SampleCount * (WorkerIndex + 1)) / WorkerCount;
            const double WorkerStart = FPlatformTime::Seconds();
            double WorkerCpuStart = 0.0;
            const bool bHaveWorkerCpuStart = GetCurrentThreadCpuSeconds(WorkerCpuStart);
            for (int64 OutputIndex = Begin; OutputIndex < End; ++OutputIndex)
            {
                const int64 LatticeIndex = bFullRing
                    ? OutputIndex : (*SampleIndices)[static_cast<int32>(OutputIndex)];
                if (LatticeIndex < 0 || LatticeIndex >= TotalLatticeSamples)
                {
                    ++WorkerNonFinite[WorkerIndex];
                    continue;
                }
                const int32 X = static_cast<int32>(LatticeIndex % Dimension);
                const int32 Y = static_cast<int32>((LatticeIndex / Dimension) % Dimension);
                const int32 Z = static_cast<int32>(LatticeIndex
                    / (static_cast<int64>(Dimension) * Dimension));
                const float Density = Generator.GetDensityAt(
                    static_cast<float>(OriginX + X),
                    static_cast<float>(OriginY + Y),
                    static_cast<float>(OriginZ + Z));
                if (OutDensities != nullptr)
                {
                    (*OutDensities)[static_cast<int32>(OutputIndex)] = Density;
                }
            }
            WorkerSeconds[WorkerIndex] = FPlatformTime::Seconds() - WorkerStart;
            double WorkerCpuEnd = 0.0;
            if (bHaveWorkerCpuStart && GetCurrentThreadCpuSeconds(WorkerCpuEnd))
            {
                WorkerCpuSeconds[WorkerIndex] = FMath::Max(0.0, WorkerCpuEnd - WorkerCpuStart);
                WorkerCpuFromOsThreadTimes[WorkerIndex] = true;
            }
            else
            {
                WorkerCpuSeconds[WorkerIndex] = WorkerSeconds[WorkerIndex];
            }
            VoxelCaveMorphology::SetResearchBrickBinningEnabled(false);
            VoxelCaveMorphology::SetResearchBrickRoomReachScale(1.0f);
        }, EParallelForFlags::None);

        OutTiming.WallSeconds = FPlatformTime::Seconds() - StartSeconds;
        OutTiming.SampleCount = SampleCount;
        for (int32 WorkerIndex = 0; WorkerIndex < WorkerCount; ++WorkerIndex)
        {
            OutTiming.WorkerCpuSeconds += WorkerCpuSeconds[WorkerIndex];
            bAllWorkerCpuFromOsThreadTimes &= WorkerCpuFromOsThreadTimes[WorkerIndex];
            OutTiming.NonFiniteSamples += WorkerNonFinite[WorkerIndex];
            bool bAlreadyCounted = false;
            for (int32 PriorWorker = 0; PriorWorker < WorkerIndex; ++PriorWorker)
            {
                bAlreadyCounted |= WorkerThreadIds[PriorWorker] == WorkerThreadIds[WorkerIndex];
            }
            OutTiming.ObservedWorkerThreads += bAlreadyCounted ? 0 : 1;
        }
        OutTiming.bWorkerCpuFromOsThreadTimes = bAllWorkerCpuFromOsThreadTimes;
        OutTiming.NanosecondsPerSample = OutTiming.SampleCount > 0
            ? OutTiming.WorkerCpuSeconds * 1.0e9 / static_cast<double>(OutTiming.SampleCount)
            : 0.0;
        if (OutDensities != nullptr)
        {
            for (const float Density : *OutDensities)
            {
                bool bSaturated = false;
                bool bFinite = false;
                QuantizeDensity(Density, bSaturated, bFinite);
                OutTiming.SaturatedSamples += bSaturated ? 1u : 0u;
                OutTiming.NonFiniteSamples += bFinite ? 0u : 1u;
            }
            OutTiming.QuantizedHash = HashQuantized(*OutDensities);
            OutTiming.RawFloatHash = HashRawFloats(*OutDensities);
        }
        return OutTiming.NonFiniteSamples == 0;
    }

    FCandidateVisits CountCandidates(
        const UVoxelGenerator& Generator,
        FIntVector OriginVoxels,
        int32 Dimension,
        const TArray<int64>& SampleIndices,
        bool bPrepared,
        float ReachScale)
    {
        FCandidateVisits Result;
        Result.SampleCount = SampleIndices.Num();
        VoxelDensityProfile::Reset();
        VoxelDensityProfile::SetMode(
            VoxelDensityProfile::EMode::Attribution,
            VoxelDensityProfile::DefaultSampleInterval);
        FTiming IgnoredTiming;
        const bool bEvaluated = Evaluate(
            Generator, OriginVoxels, Dimension, &SampleIndices, false,
            bPrepared, ReachScale, nullptr, IgnoredTiming);
        const VoxelDensityProfile::FSnapshot Snapshot = VoxelDensityProfile::Snapshot();
        VoxelDensityProfile::SetMode(
            VoxelDensityProfile::EMode::Disabled,
            VoxelDensityProfile::DefaultSampleInterval);
        VoxelDensityProfile::Reset();
        if (!bEvaluated)
        {
            return Result;
        }
        using VoxelDensityProfile::ECounter;
        Result.Room = Snapshot.Counters[static_cast<int32>(ECounter::CaveRoomCandidates)];
        Result.Tunnel = Snapshot.Counters[static_cast<int32>(ECounter::CaveTunnelCandidates)];
        Result.TunnelCore = Snapshot.Counters[static_cast<int32>(ECounter::TunnelCoreCandidates)];
        Result.Passage = Snapshot.Counters[static_cast<int32>(ECounter::PassageCandidates)];
        Result.TunnelCacheLookups = Snapshot.Counters[static_cast<int32>(ECounter::TunnelCacheLookup)];
        Result.TunnelCacheHits = Snapshot.Counters[static_cast<int32>(ECounter::TunnelCacheHit)];
        Result.TunnelCacheMisses = Snapshot.Counters[static_cast<int32>(ECounter::TunnelCacheMiss)];
        Result.SdfCacheBuilds = Snapshot.Counters[static_cast<int32>(ECounter::SdfCacheBuild)];
        Result.Total = Result.Room + Result.Tunnel + Result.TunnelCore + Result.Passage;
        Result.PerSample = Result.SampleCount > 0
            ? static_cast<double>(Result.Total) / Result.SampleCount : 0.0;
        Result.AttributionSampleInterval = Snapshot.SampleInterval;
        const auto AttributedNanosecondsPerSample = [&](VoxelDensityProfile::EBucket Bucket)
        {
            return Result.SampleCount > 0
                ? FPlatformTime::ToSeconds64(Snapshot.Cycles[static_cast<int32>(Bucket)])
                    * 1.0e9 / static_cast<double>(Result.SampleCount)
                : 0.0;
        };
        Result.GetDensityAtAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::GetDensityAt);
        Result.DensityPrologueAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityPrologue);
        Result.DensityCoreAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityCore);
        Result.DensityDisturbancesAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityDisturbances);
        Result.DensityStructuralPostsAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityStructuralPosts);
        Result.DensityBoundarySealAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityBoundarySeal);
        Result.DensityDiffLayerAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityDiffLayer);
        Result.DensityTailAttributedNsPerSample = AttributedNanosecondsPerSample(
            VoxelDensityProfile::EBucket::DensityTail);
        return Result;
    }
}