#include "VoxelForgeRegionEvalExperiment.h"

#include "Async/ParallelFor.h"
#include "HAL/PlatformTime.h"
#include "VoxelCaveMorphology.h"
#include "VoxelDensityProfile.h"
#include "VoxelGenerator.h"
#include "VoxelStackSampler.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;
        }

        struct FRegionBrickSample
        {
            int32 LocalIndex = 0;
            int32 OutputIndex = 0;
        };

        struct FRegionBrick
        {
            FIntVector Origin = FIntVector::ZeroValue;
            int32 SizeX = 0;
            int32 SizeY = 0;
            int32 SizeZ = 0;
            int32 LatticeX = 0;
            int32 LatticeY = 0;
            int32 LatticeZ = 0;
            TArray<FRegionBrickSample> RequestedSamples;
        };
    }

    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,
        bool bBrickProgram,
        int32 RequestedWorkerCount)
    {
        OutTiming = FTiming();
        const int32 NumWorkers = FMath::Clamp(RequestedWorkerCount, 1, 32);
        OutTiming.Workers = NumWorkers;
        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();
        TArray<FRegionBrick> Bricks;
        TArray<int32> WorkBricks;
        if (bBrickProgram)
        {
            const int32 LastX = OriginVoxels.X + Dimension - 1;
            const int32 LastY = OriginVoxels.Y + Dimension - 1;
            const int32 LastZ = OriginVoxels.Z + Dimension - 1;
            const int32 FirstCellX = FMath::FloorToInt(static_cast<float>(OriginVoxels.X) / 8.0f);
            const int32 FirstCellY = FMath::FloorToInt(static_cast<float>(OriginVoxels.Y) / 8.0f);
            const int32 FirstCellZ = FMath::FloorToInt(static_cast<float>(OriginVoxels.Z) / 8.0f);
            const int32 LastCellX = FMath::FloorToInt(static_cast<float>(LastX) / 8.0f);
            const int32 LastCellY = FMath::FloorToInt(static_cast<float>(LastY) / 8.0f);
            const int32 LastCellZ = FMath::FloorToInt(static_cast<float>(LastZ) / 8.0f);
            const int32 CellsX = LastCellX - FirstCellX + 1;
            const int32 CellsY = LastCellY - FirstCellY + 1;
            const int32 CellsZ = LastCellZ - FirstCellZ + 1;
            Bricks.Reserve(CellsX * CellsY * CellsZ);
            for (int32 CellZ = FirstCellZ; CellZ <= LastCellZ; ++CellZ)
            for (int32 CellY = FirstCellY; CellY <= LastCellY; ++CellY)
            for (int32 CellX = FirstCellX; CellX <= LastCellX; ++CellX)
            {
                FRegionBrick& Brick = Bricks.AddDefaulted_GetRef();
                const int32 CellMinX = CellX * 8;
                const int32 CellMinY = CellY * 8;
                const int32 CellMinZ = CellZ * 8;
                Brick.Origin = FIntVector(
                    FMath::Max(CellMinX, OriginVoxels.X),
                    FMath::Max(CellMinY, OriginVoxels.Y),
                    FMath::Max(CellMinZ, OriginVoxels.Z));
                const int32 EndX = FMath::Min(CellMinX + 8, LastX + 1);
                const int32 EndY = FMath::Min(CellMinY + 8, LastY + 1);
                const int32 EndZ = FMath::Min(CellMinZ + 8, LastZ + 1);
                Brick.SizeX = EndX - Brick.Origin.X;
                Brick.SizeY = EndY - Brick.Origin.Y;
                Brick.SizeZ = EndZ - Brick.Origin.Z;
                Brick.LatticeX = Brick.Origin.X - OriginVoxels.X;
                Brick.LatticeY = Brick.Origin.Y - OriginVoxels.Y;
                Brick.LatticeZ = Brick.Origin.Z - OriginVoxels.Z;
            }

            if (bFullRing)
            {
                WorkBricks.Reserve(Bricks.Num());
                for (int32 BrickIndex = 0; BrickIndex < Bricks.Num(); ++BrickIndex)
                {
                    WorkBricks.Add(BrickIndex);
                }
            }
            else
            {
                WorkBricks.Reserve(Bricks.Num());
                for (int32 OutputIndex = 0; OutputIndex < SampleIndices->Num(); ++OutputIndex)
                {
                    const int64 LatticeIndex = (*SampleIndices)[OutputIndex];
                    if (LatticeIndex < 0 || LatticeIndex >= TotalLatticeSamples)
                    {
                        return false;
                    }
                    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 int32 WorldX = OriginVoxels.X + X;
                    const int32 WorldY = OriginVoxels.Y + Y;
                    const int32 WorldZ = OriginVoxels.Z + Z;
                    const int32 CellX = FMath::FloorToInt(static_cast<float>(WorldX) / 8.0f);
                    const int32 CellY = FMath::FloorToInt(static_cast<float>(WorldY) / 8.0f);
                    const int32 CellZ = FMath::FloorToInt(static_cast<float>(WorldZ) / 8.0f);
                    const int32 BrickIndex = ((CellZ - FirstCellZ) * CellsY
                        + (CellY - FirstCellY)) * CellsX + (CellX - FirstCellX);
                    FRegionBrick& Brick = Bricks[BrickIndex];
                    const int32 LocalX = WorldX - Brick.Origin.X;
                    const int32 LocalY = WorldY - Brick.Origin.Y;
                    const int32 LocalZ = WorldZ - Brick.Origin.Z;
                    FRegionBrickSample& BrickSample = Brick.RequestedSamples.AddDefaulted_GetRef();
                    BrickSample.LocalIndex =
                        (LocalZ * Brick.SizeY + LocalY) * Brick.SizeX + LocalX;
                    BrickSample.OutputIndex = OutputIndex;
                }
                for (int32 BrickIndex = 0; BrickIndex < Bricks.Num(); ++BrickIndex)
                {
                    if (!Bricks[BrickIndex].RequestedSamples.IsEmpty())
                    {
                        WorkBricks.Add(BrickIndex);
                    }
                }
            }
        }
        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;
        TArray<double> WorkerSeconds;
        TArray<double> WorkerCpuSeconds;
        TArray<uint8> WorkerCpuFromOsThreadTimes;
        TArray<uint32> WorkerThreadIds;
        TArray<uint64> WorkerNonFinite;
        TArray<uint64> WorkerEvaluatedSamples;
        TArray<uint64> WorkerFallbackSamples;
        TArray<uint64> WorkerEvaluationFailures;
        WorkerSeconds.SetNumZeroed(NumWorkers);
        WorkerCpuSeconds.SetNumZeroed(NumWorkers);
        WorkerCpuFromOsThreadTimes.SetNumZeroed(NumWorkers);
        WorkerThreadIds.SetNumZeroed(NumWorkers);
        WorkerNonFinite.SetNumZeroed(NumWorkers);
        WorkerEvaluatedSamples.SetNumZeroed(NumWorkers);
        WorkerFallbackSamples.SetNumZeroed(NumWorkers);
        WorkerEvaluationFailures.SetNumZeroed(NumWorkers);
        bool bAllWorkerCpuFromOsThreadTimes = true;
        ParallelFor(NumWorkers, [&](int32 WorkerIndex)
        {
            // Region-eval workers are not launched through AVoxelWorld::LoadTile, so they do not
            // inherit its task-level sampler registration. Keep this registration scoped to the
            // research worker partition; with sampling disabled it is one relaxed flag load.
            FScopedVoxelStackRegistration StackRegistration;
            VoxelCaveMorphology::SetResearchBrickBinningEnabled(bPrepared);
            VoxelCaveMorphology::SetResearchBrickRoomReachScale(ReachScale);
            WorkerThreadIds[WorkerIndex] = FPlatformTLS::GetCurrentThreadId();
            const int64 WorkCount = bBrickProgram ? WorkBricks.Num() : SampleCount;
            const int64 Begin = (WorkCount * WorkerIndex) / NumWorkers;
            const int64 End = (WorkCount * (WorkerIndex + 1)) / NumWorkers;
            const double WorkerStart = FPlatformTime::Seconds();
            double WorkerCpuStart = 0.0;
            const bool bHaveWorkerCpuStart = GetCurrentThreadCpuSeconds(WorkerCpuStart);
            TArray<float> BrickDensities;
            for (int64 OutputIndex = Begin; OutputIndex < End; ++OutputIndex)
            {
                if (bBrickProgram)
                {
                    const FRegionBrick& Brick = Bricks[WorkBricks[static_cast<int32>(OutputIndex)]];
                    int32 BrickFallbackSamples = 0;
                    if (!Generator.EvaluateDensityBrick(
                            Brick.Origin, Brick.SizeX, Brick.SizeY, Brick.SizeZ,
                            BrickDensities, BrickFallbackSamples)
                        || BrickDensities.Num() != Brick.SizeX * Brick.SizeY * Brick.SizeZ)
                    {
                        ++WorkerEvaluationFailures[WorkerIndex];
                        continue;
                    }
                    const int32 BrickSampleCount = BrickDensities.Num();
                    WorkerEvaluatedSamples[WorkerIndex] += static_cast<uint64>(BrickSampleCount);
                    WorkerFallbackSamples[WorkerIndex] += static_cast<uint64>(BrickFallbackSamples);
                    if (bFullRing && OutDensities != nullptr)
                    {
                        for (int32 Bz = 0; Bz < Brick.SizeZ; ++Bz)
                        for (int32 By = 0; By < Brick.SizeY; ++By)
                        for (int32 Bx = 0; Bx < Brick.SizeX; ++Bx)
                        {
                            const int32 LocalIndex = (Bz * Brick.SizeY + By) * Brick.SizeX + Bx;
                            const int64 LatticeIndex = LinearIndex(
                                Brick.LatticeX + Bx, Brick.LatticeY + By,
                                Brick.LatticeZ + Bz, Dimension);
                            (*OutDensities)[static_cast<int32>(LatticeIndex)] = BrickDensities[LocalIndex];
                        }
                    }
                    else if (OutDensities != nullptr)
                    {
                        for (const FRegionBrickSample& Sample : Brick.RequestedSamples)
                        {
                            (*OutDensities)[Sample.OutputIndex] = BrickDensities[Sample.LocalIndex];
                        }
                    }
                    continue;
                }
                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;
        OutTiming.bBrickProgram = bBrickProgram;
        OutTiming.BrickCount = bBrickProgram ? WorkBricks.Num() : 0;
        for (int32 WorkerIndex = 0; WorkerIndex < NumWorkers; ++WorkerIndex)
        {
            OutTiming.WorkerCpuSeconds += WorkerCpuSeconds[WorkerIndex];
            bAllWorkerCpuFromOsThreadTimes &= WorkerCpuFromOsThreadTimes[WorkerIndex] != 0;
            OutTiming.NonFiniteSamples += WorkerNonFinite[WorkerIndex];
            OutTiming.EvaluatedSamples += bBrickProgram
                ? static_cast<int64>(WorkerEvaluatedSamples[WorkerIndex]) : 0;
            OutTiming.FallbackSamples += static_cast<int64>(WorkerFallbackSamples[WorkerIndex]);
            if (WorkerEvaluationFailures[WorkerIndex] != 0)
            {
                return false;
            }
            bool bAlreadyCounted = false;
            for (int32 PriorWorker = 0; PriorWorker < WorkerIndex; ++PriorWorker)
            {
                bAlreadyCounted |= WorkerThreadIds[PriorWorker] == WorkerThreadIds[WorkerIndex];
            }
            OutTiming.ObservedWorkerThreads += bAlreadyCounted ? 0 : 1;
        }
        OutTiming.bWorkerCpuFromOsThreadTimes = bAllWorkerCpuFromOsThreadTimes;
        if (OutTiming.EvaluatedSamples > 0)
        {
            OutTiming.FallbackShare = static_cast<double>(OutTiming.FallbackSamples)
                / static_cast<double>(OutTiming.EvaluatedSamples);
        }
        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;
    }
}