/* Copyright UNmisterIZE(Sebastien Durocher) 2025 All Rights Reserved.

Project: NoiseMapGenerator Plugin
File: NoiseMapGenerator.cpp
*/

#include "NoiseMapGenerator.h"
#include "Engine/Texture2D.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "Misc/PackageName.h"
#include "UObject/Package.h"
#include "UObject/SavePackage.h"

////// ---------------------------------------------------------
namespace TileableNoise
{
    static FORCEINLINE float Fade(float t){ return t*t*t*(t*(t*6 - 15) + 10); }
    static FORCEINLINE float Lerp(float a,float b,float t){ return a + t*(b - a); }

    // Permutation table (512 entries = 256 duplicated)
    struct Perm
    {
        TArray<int32> P;
        explicit Perm(int32 Seed)
        {
            P.SetNumUninitialized(512);
            TArray<int32> base; base.SetNumUninitialized(256);
            for (int32 i=0;i<256;++i) base[i]=i;
            FRandomStream RNG(Seed);
            for (int32 i=255;i>0;--i){ const int32 j=RNG.RandRange(0,i); Swap(base[i],base[j]); }
            for (int32 i=0;i<512;++i) P[i]=base[i & 255];
        }
        FORCEINLINE int32 H4(int32 x,int32 y,int32 z,int32 w) const
        { return P[P[P[P[x & 255] + (y & 255)] + (z & 255)] + (w & 255)]; }
    };

    // 4D gradient (32 dirs)
    static FORCEINLINE float Grad4(int32 h, float x,float y,float z,float w)
    {
        const int32 b = h & 31;
        switch (b)
        {
            case  0: return  x + y + z + w;  case  1: return  x + y + z - w;
            case  2: return  x + y - z + w;  case  3: return  x + y - z - w;
            case  4: return  x - y + z + w;  case  5: return  x - y + z - w;
            case  6: return  x - y - z + w;  case  7: return  x - y - z - w;
            case  8: return -x + y + z + w;  case  9: return -x + y + z - w;
            case 10: return -x + y - z + w;  case 11: return -x + y - z - w;
            case 12: return -x - y + z + w;  case 13: return -x - y + z - w;
            case 14: return -x - y - z + w;  case 15: return -x - y - z - w;
            case 16: return  x + y + z;      case 17: return  x + y - z;
            case 18: return  x - y + z;      case 19: return  x - y - z;
            case 20: return -x + y + z;      case 21: return -x + y - z;
            case 22: return -x - y + z;      case 23: return -x - y - z;
            case 24: return  x + z + w;      case 25: return  x + z - w;
            case 26: return  x - z + w;      case 27: return  x - z - w;
            case 28: return -x + z + w;      case 29: return -x + z - w;
            case 30: return -x - z + w;      default:  return -x - z - w;
        }
    }

    static float Perlin4D(float x,float y,float z,float w,const Perm& perm)
    {
        const int32 X0 = FMath::FloorToInt(x), Y0 = FMath::FloorToInt(y), Z0 = FMath::FloorToInt(z), W0 = FMath::FloorToInt(w);
        const float xf = x - X0, yf = y - Y0, zf = z - Z0, wf = w - W0;
        const float u = Fade(xf), v = Fade(yf), s = Fade(zf), t = Fade(wf);
        const int32 X1=X0+1,Y1=Y0+1,Z1=Z0+1,W1=W0+1;

        const float n0000 = Grad4(perm.H4(X0,Y0,Z0,W0), xf,     yf,     zf,     wf);
        const float n1000 = Grad4(perm.H4(X1,Y0,Z0,W0), xf-1.f, yf,     zf,     wf);
        const float n0100 = Grad4(perm.H4(X0,Y1,Z0,W0), xf,     yf-1.f, zf,     wf);
        const float n1100 = Grad4(perm.H4(X1,Y1,Z0,W0), xf-1.f, yf-1.f, zf,     wf);
        const float n0010 = Grad4(perm.H4(X0,Y0,Z1,W0), xf,     yf,     zf-1.f, wf);
        const float n1010 = Grad4(perm.H4(X1,Y0,Z1,W0), xf-1.f, yf,     zf-1.f, wf);
        const float n0110 = Grad4(perm.H4(X0,Y1,Z1,W0), xf,     yf-1.f, zf-1.f, wf);
        const float n1110 = Grad4(perm.H4(X1,Y1,Z1,W0), xf-1.f, yf-1.f, zf-1.f, wf);

        const float n0001 = Grad4(perm.H4(X0,Y0,Z0,W1), xf,     yf,     zf,     wf-1.f);
        const float n1001 = Grad4(perm.H4(X1,Y0,Z0,W1), xf-1.f, yf,     zf,     wf-1.f);
        const float n0101 = Grad4(perm.H4(X0,Y1,Z0,W1), xf,     yf-1.f, zf,     wf-1.f);
        const float n1101 = Grad4(perm.H4(X1,Y1,Z0,W1), xf-1.f, yf-1.f, zf,     wf-1.f);
        const float n0011 = Grad4(perm.H4(X0,Y0,Z1,W1), xf,     yf,     zf-1.f, wf-1.f);
        const float n1011 = Grad4(perm.H4(X1,Y0,Z1,W1), xf-1.f, yf,     zf-1.f, wf-1.f);
        const float n0111 = Grad4(perm.H4(X0,Y1,Z1,W1), xf,     yf-1.f, zf-1.f, wf-1.f);
        const float n1111 = Grad4(perm.H4(X1,Y1,Z1,W1), xf-1.f, yf-1.f, zf-1.f, wf-1.f);

        const float nx00  = Lerp(n0000, n1000, u);
        const float nx10  = Lerp(n0100, n1100, u);
        const float nx01  = Lerp(n0010, n1010, u);
        const float nx11  = Lerp(n0110, n1110, u);
        const float nxy0  = Lerp(nx00,  nx10,  v);
        const float nxy1  = Lerp(nx01,  nx11,  v);
        const float nxyz0 = Lerp(nxy0,  nxy1,  s);

        const float nx00b  = Lerp(n0001, n1001, u);
        const float nx10b  = Lerp(n0101, n1101, u);
        const float nx01b  = Lerp(n0011, n1011, u);
        const float nx11b  = Lerp(n0111, n1111, u);
        const float nxy0b  = Lerp(nx00b, nx10b, v);
        const float nxy1b  = Lerp(nx01b, nx11b, v);
        const float nxyz1  = Lerp(nxy0b, nxy1b, s);

        return Lerp(nxyz0, nxyz1, t); // [-1,1]
    }

    static float FBM4D(float x,float y,float z,float w,const Perm& perm,int32 Octaves,float Lac,float Gain)
    {
        float amp=1.f,sum=0.f,range=0.f;
        float X=x,Y=y,Z=z,W=w;
        for (int32 o=0;o<Octaves;++o)
        {
            sum += amp * Perlin4D(X,Y,Z,W,perm);
            range += amp;
            X*=Lac; Y*=Lac; Z*=Lac; W*=Lac;
            amp*=Gain;
        }
        return (range>0.f)? sum/range : 0.f;
    }

    // Hash utilities for Worley
    static FORCEINLINE uint32 Hash2D(int32 x,int32 y,uint32 seed)
    {
        uint32 h = (uint32)x * 0x27d4eb2d ^ (uint32)y * 0x165667b1u ^ seed*0x9e3779b9u;
        h ^= (h >> 16); h *= 0x7feb352d; h ^= (h >> 15); h *= 0x846ca68b; h ^= (h >> 16);
        return h;
    }
    static FORCEINLINE float Frand01(uint32& state)
    {
        state ^= state << 13; state ^= state >> 17; state ^= state << 5;
        return (state & 0x00FFFFFF) / 16777216.0f;
    }

    // Torus-safe distance between (u,v) and feature in neighbor cell (iu,iv) with wrap
    static float WorleyF1_Torus(float u,float v,int32 CellsU,int32 CellsV,uint32 seed)
    {
        // Base integer cell
        const float fu = u * CellsU;
        const float fv = v * CellsV;
        const int32 iu0 = FMath::FloorToInt(fu);
        const int32 iv0 = FMath::FloorToInt(fv);

        float best = 1e9f;

        for (int du=-1; du<=1; ++du)
        for (int dv=-1; dv<=1; ++dv)
        {
            const int32 iu = (iu0 + du + CellsU) % CellsU;
            const int32 iv = (iv0 + dv + CellsV) % CellsV;

            uint32 st = Hash2D(iu, iv, seed);
            const float jitterX = Frand01(st);
            const float jitterY = Frand01(st);

            const float fx = (float(iu) + jitterX) / float(CellsU);
            const float fy = (float(iv) + jitterY) / float(CellsV);

            // torus wrap distance in [0,1]
            float dx = FMath::Abs(u - fx); dx = FMath::Min(dx, 1.0f - dx);
            float dy = FMath::Abs(v - fy); dy = FMath::Min(dy, 1.0f - dy);

            best = FMath::Min(best, FMath::Sqrt(dx*dx + dy*dy));
        }
        return best; // [0..~0.5]
    }

}
/////// ---------------------------------------------------------

static void MakeUniqueAssetPath(const FString& Folder, const FString& BaseName, bool bOverwrite, FString& OutPkg, FString& OutAsset)
{
    FString CleanFolder = Folder;
    if (!CleanFolder.StartsWith(TEXT("/"))) CleanFolder = TEXT("/") + CleanFolder;
    if (!CleanFolder.StartsWith(TEXT("/Game"))) CleanFolder = TEXT("/Game") + CleanFolder;

    FString CandidateName = BaseName;
    FString CandidatePkg  = CleanFolder / CandidateName;

    if (bOverwrite)
    {
        OutPkg = CandidatePkg;
        OutAsset = CandidateName;
        return;
    }

    int32 Suffix = 1;
    while (FPackageName::DoesPackageExist(CandidatePkg))
    {
        CandidateName = FString::Printf(TEXT("%s_%03d"), *BaseName, Suffix++);
        CandidatePkg  = CleanFolder / CandidateName;
    }
    OutPkg = CandidatePkg;
    OutAsset = CandidateName;
}


bool FNoiseMapGenerator::GeneratePixels(const FNoiseBakeSettings& S, int32 RepeatsUV, TArray<FColor>& OutPixels, int32& OutW, int32& OutH)
{
    // --- size ---
    const int32 W = S.Width;
    const int32 H = S.Height;
    if (W <= 0 || H <= 0)
    {
        OutW = OutH = 0;
        OutPixels.Reset();
        return false;
    }
    OutW = W;
    OutH = H;
    OutPixels.SetNumUninitialized(W * H);

    // --- common setup ---
    const float invW  = 1.f / float(W);
    const float invH  = 1.f / float(H);
    const float TwoPi = 6.28318530717958647692f;
    const float Scale = 1.0f / FMath::Max(S.Frequency, 1e-6f);

    TileableNoise::Perm perm(S.Seed);

    // 4D torus mapping (period-1 in u,v)
    auto AnglesFromUV = [&](float u, float v)
    {
        const float th = TwoPi * u;
        const float ph = TwoPi * v;
        struct { float X, Y, Z, W; } R
        {
            FMath::Cos(th) * Scale,
            FMath::Sin(th) * Scale,
            FMath::Cos(ph) * Scale,
            FMath::Sin(ph) * Scale
        };
        return R;
    };
    
    // Base sampler per algorithm, returns [-1,1]
    auto SampleBase = [&](float u, float v) -> float
    {
        switch (S.Algorithm)
        {
        case ENoiseAlgorithm::Perlin:
            {
                const auto C = AnglesFromUV(u, v);
                return (S.Mode == ENoiseMode::FBM)
                    ? TileableNoise::FBM4D(C.X, C.Y, C.Z, C.W, perm, S.Octaves, S.Lacunarity, S.Gain)
                    : TileableNoise::Perlin4D(C.X, C.Y, C.Z, C.W, perm);
            }

        case ENoiseAlgorithm::RidgedFBM:
            {
                float amp     = 0.5f;
                float sum     = 0.f;
                float range   = 0.f;
                float freqMul = 1.f;

                for (int32 o = 0; o < S.Octaves; ++o)
                {
                    const float ou = FMath::Frac(u * freqMul);
                    const float ov = FMath::Frac(v * freqMul);
                    const auto C   = AnglesFromUV(ou, ov);
                    const float p  = TileableNoise::Perlin4D(C.X, C.Y, C.Z, C.W, perm); // [-1,1]
                    const float r  = 1.f - FMath::Abs(p); // [0,1]

                    sum   += amp * r;
                    range += amp;

                    freqMul *= S.Lacunarity;
                    amp     *= S.Gain;
                }

                return (range > 0.f) ? (sum / range) * 2.f - 1.f : 0.f;
            }

        case ENoiseAlgorithm::Worley:
            {
                const float cellsF = FMath::Clamp(
                    (1.0f / FMath::Max(S.Frequency, 1e-6f)) * 0.75f,
                    1.f, 512.f
                );

                const int32 CellsU = FMath::Max(1, int32(FMath::RoundToInt(cellsF)) * RepeatsUV);
                const int32 CellsV = FMath::Max(1, int32(FMath::RoundToInt(cellsF)) * RepeatsUV);

                const float d   = TileableNoise::WorleyF1_Torus(u, v, CellsU, CellsV, (uint32)S.Seed); // [0..~0.5]
                const float v01 = 1.0f - FMath::Clamp(d * 2.0f, 0.f, 1.f);
                return v01 * 2.f - 1.f;
            }
        default:
            return 0.f;
        }
    };
    
    auto SampleNthOctave01 = [&](float u, float v, int32 OctIdx) -> float
    {
        switch (S.Algorithm)
        {
        case ENoiseAlgorithm::Perlin:
            {
                // Higher octave = higher frequency
                const float freqMul = FMath::Pow(S.Lacunarity, OctIdx);
                const auto  C       = AnglesFromUV(u * freqMul, v * freqMul);
                const float p       = TileableNoise::Perlin4D(C.X, C.Y, C.Z, C.W, perm); // [-1,1]
                return 0.5f * (p + 1.f); // → [0,1]
            }

        case ENoiseAlgorithm::RidgedFBM:
            {
                const float freqMul = FMath::Pow(S.Lacunarity, OctIdx);
                const float ou      = FMath::Frac(u * freqMul);
                const float ov      = FMath::Frac(v * freqMul);
                const auto  C       = AnglesFromUV(ou, ov);
                const float p       = TileableNoise::Perlin4D(C.X, C.Y, C.Z, C.W, perm); // [-1,1]
                const float r       = 1.f - FMath::Abs(p); // [0,1] – "ridge" shape
                return r;
            }

        case ENoiseAlgorithm::Worley:
            {
                // More cells per octave for finer detail
                const float cellsFBase = FMath::Clamp(
                    (1.0f / FMath::Max(S.Frequency, 1e-6f)) * 0.75f,
                    1.f, 512.f
                );
                const float octaveMul = FMath::Pow(S.Lacunarity, OctIdx);
                const int32 CellsU    = FMath::Max(1, int32(FMath::RoundToInt(cellsFBase * octaveMul)) * RepeatsUV);
                const int32 CellsV    = FMath::Max(1, int32(FMath::RoundToInt(cellsFBase * octaveMul)) * RepeatsUV);

                const float d   = TileableNoise::WorleyF1_Torus(u, v, CellsU, CellsV, (uint32)S.Seed);
                const float v01 = 1.0f - FMath::Clamp(d * 2.0f, 0.f, 1.f); // [0,1], "cells" as bright
                return v01;
            }

        default:
            // Fallback mid-gray
                return 0.5f;
        }
    };
    
    // Write pixels
    if (S.OutputMode == ENoiseOutputMode::Grayscale)
    {
        float minV =  1e9f;
        float maxV = -1e9f;
        TArray<float> tmp;
        tmp.SetNumUninitialized(W * H);

        int32 I = 0;
        for (int32 y = 0; y < H; ++y)
        {
            for (int32 x = 0; x < W; ++x)
            {
                const float u = float(RepeatsUV) * float(x) * invW; 
                const float v = float(RepeatsUV) * float(y) * invH; 

                const float vraw = SampleBase(u, v); // [-1,1]
                tmp[I++] = vraw;
                minV = FMath::Min(minV, vraw);
                maxV = FMath::Max(maxV, vraw);
            }
        }

        const float invRange =
            (S.bNormalize && maxV > minV) ? 1.f / (maxV - minV) : 0.5f;

        I = 0;
        for (int32 y = 0; y < H; ++y)
        {
            for (int32 x = 0; x < W; ++x)
            {
                float v = tmp[I++];
                v = S.bNormalize
                    ? (v - minV) * invRange       // [0,1]
                    : 0.5f * (v + 1.f);           // [-1,1] → [0,1]

                const uint8 g = (uint8)FMath::Clamp(
                    FMath::RoundToInt(v * 255.f), 0, 255);

                OutPixels[y * W + x] = FColor(g, g, g, 255);
            }
        }
    }
    else // Split First 3 Octaves to RGB channles
    {
        for (int32 y = 0; y < H; ++y)
        {
            for (int32 x = 0; x < W; ++x)
            {
                const float u = float(RepeatsUV) * float(x) * invW;
                const float v = float(RepeatsUV) * float(y) * invH;

                const float r = SampleNthOctave01(u, v, 0);
                const float g = SampleNthOctave01(u, v, 1);
                const float b = SampleNthOctave01(u, v, 2);

                OutPixels[y * W + x] = FColor(
                    (uint8)FMath::Clamp(FMath::RoundToInt(r * 255.f), 0, 255),
                    (uint8)FMath::Clamp(FMath::RoundToInt(g * 255.f), 0, 255),
                    (uint8)FMath::Clamp(FMath::RoundToInt(b * 255.f), 0, 255),
                    255);
            }
        }
    }

    return true;
}

void FNoiseMapGenerator::BakeFromPixels(
    const FNoiseBakeSettings& S,
    const TArray<FColor>& Pixels,
    int32 W,
    int32 H)
{
    if (W <= 0 || H <= 0 || Pixels.Num() != W * H)
    {
        return;
    }

    // --- Resolve package path/name ---
    FString OutPath = S.OutputPath.IsEmpty()
        ? TEXT("/Game/Noise")
        : S.OutputPath;

    if (!OutPath.StartsWith(TEXT("/Game")))
    {
        OutPath = TEXT("/Game/Noise");
    }

    FString PackageName, AssetName;
    MakeUniqueAssetPath(OutPath, S.BaseName, S.bOverwrite, PackageName, AssetName);

    UPackage* Package = CreatePackage(*PackageName);
    if (!Package)
    {
        return;
    }

    UTexture2D* Tex = NewObject<UTexture2D>(
        Package,
        *AssetName,
        RF_Public | RF_Standalone
    );
    if (!Tex)
    {
        return;
    }

    Tex->Modify();

    // Match preview flags
    Tex->SRGB = false;

    if (S.OutputMode == ENoiseOutputMode::Grayscale)
    {
        Tex->CompressionSettings = TC_Grayscale;
    }
    else
    {
        Tex->CompressionSettings = TC_VectorDisplacementmap; // linear RGB
    }

    Tex->CompressionNone  = true;
    Tex->DeferCompression = false;
    Tex->MipGenSettings   = TMGS_NoMipmaps;
    Tex->AddressX         = TA_Wrap;
    Tex->AddressY         = TA_Wrap;
    Tex->LODGroup         = TEXTUREGROUP_World;

    // Copy pixels into source mip
    Tex->Source.Init(W, H, 1, 1, TSF_BGRA8);
    {
        uint8* Src = Tex->Source.LockMip(0);
        const int64 Bytes = int64(W) * H * sizeof(FColor);
        FMemory::Memcpy(Src, Pixels.GetData(), Bytes);
        Tex->Source.UnlockMip(0);
    }

    Tex->PostEditChange();
    Tex->UpdateResource();

    FAssetRegistryModule::AssetCreated(Tex);
    Package->MarkPackageDirty();

    const FString PkgFile =
        FPackageName::LongPackageNameToFilename(
            PackageName,
            FPackageName::GetAssetPackageExtension()
        );

    FSavePackageArgs Args;
    Args.TopLevelFlags       = RF_Public | RF_Standalone;
    Args.SaveFlags           = SAVE_None;
    Args.Error               = GError;
    Args.bWarnOfLongFilename = true;

    UPackage::SavePackage(Package, Tex, *PkgFile, Args);
}

void FNoiseMapGenerator::BakeNoiseTexture(const FNoiseBakeSettings& S)
{
    TArray<FColor> Pixels;
    int32 W = 0, H = 0;

    // IMPORTANT: RepeatsUV = 1 for a single seamless period
    if (!GeneratePixels(S, /*RepeatsUV=*/1, Pixels, W, H))
    {
        return;
    }

    BakeFromPixels(S, Pixels, W, H);
}
