import bindless;
import common_structs;
import ibl_common;

/*
| Axis | Layer | Up |
|:----:|:-----:|:--:|
|  +x  |   0   | -y |
|  -x  |   1   | -y |
|  +y  |   2   | +z |
|  -y  |   3   | -z |
|  -z  |   4   | -y |
|  +z  |   5   | -y |
*/

struct DiffuseBlock {
    SamplerCube<float4>.Handle skybox;
    RWTexture2DArray<float4>.Handle outputTexture;
    uint cubeSide;
};


[shader("compute")]
[numthreads(16, 16, 1)]
void diffuseIrradiance(
    uint3 GlobalInvocationID: SV_DispatchThreadID,
    uniform DiffuseBlock diffusePcb,
    )
{
    float3 Forward, Up, Right;

    Forward = GetCubeDir(GlobalInvocationID, diffusePcb.cubeSide);
    Up = abs(Forward.y) < 1.0f ? float3(0.0f, 1.0f, 0.0f) : float3(1.0f, 0.0f, 0.0f); // 0.01f offset to 
    Right = normalize(cross(Up, Forward));
    Up = normalize(cross(Forward, Right));
    
    float3 Irradiance = 0.0f.xxx;
    float3 IrrDirr = 0.0f.xxx;
    float SampleStep = 0.005f;
    float SampleCount = 0.0f;
  
    for (float Azimuth = 0.0f; Azimuth < TAU; Azimuth += SampleStep)
    {
        for (float Zenith = 0.0f; Zenith < HALF_PI; Zenith += SampleStep)
        {
            float3 DirectionTanSpace = float3(sin(Zenith) * cos(Azimuth), sin(Zenith) * sin(Azimuth), cos(Zenith));
            float3 DirectionWorld = DirectionTanSpace.x * Right + DirectionTanSpace.y * Up + DirectionTanSpace.z * Forward;

            Irradiance += diffusePcb.skybox.SampleLevel(DirectionWorld, 0).xyz * (cos(Zenith) * sin(Zenith));
            SampleCount++;
        }
    }

    diffusePcb.outputTexture[GlobalInvocationID.xyz] = PI * float4(Irradiance * (1.0f / SampleCount), 1.0f);
}

struct PrefilterBlock {
    SamplerCube<float4>.Handle skybox;
    RWTexture2DArray<float4>.Handle outputTexture;
    uint cubeSide;
    float roughness;
    uint envRes;
};

[shader("compute")]
[numthreads(16, 16, 1)]
void prefilter(uint3 GlobalInvocationID: SV_DispatchThreadID, uniform PrefilterBlock prefilterPcb)
{
    float3 Normal = GetCubeDir(GlobalInvocationID, prefilterPcb.cubeSide);
    float3 ViewDir = Normal;

    const uint SAMPLE_COUNT = 2048u;
    float TotalWeight = 0.0f;
    float3 PrefilterColor = 0.0f.xxx;
    for (uint i = 0u; i < SAMPLE_COUNT; ++i)
    {
        float2 Xi = Hammersley(i, SAMPLE_COUNT);
        float3 Halfway = ImportanceSampleGGX(Xi, Normal, prefilterPcb.roughness);
        float3 LightDir = normalize(2.0 * dot(ViewDir, Halfway) * Halfway - ViewDir);

        float NdotH = max(dot(Normal, Halfway), 0.0f);

        float MipLevel = GetSampleMipLevel(NdotH, NdotH /* N = V :: NdotH = HdotV */, SAMPLE_COUNT, prefilterPcb.roughness, prefilterPcb.envRes);

        float NdotL = max(dot(Normal, LightDir), 0.0);
        if (NdotL > 0.0)
        {
            PrefilterColor += prefilterPcb.skybox.SampleLevel(LightDir, MipLevel).rgb * NdotL;
            TotalWeight += NdotL;
        }
    }
    PrefilterColor = PrefilterColor / TotalWeight;

    prefilterPcb.outputTexture[GlobalInvocationID] = float4(PrefilterColor, 1.0f);
}

struct BrdfBlock {
    RWTexture2D<float2>.Handle outputTexture;
};

[shader("compute")]
[numthreads(16, 16, 1)]
void brdfLut(
    uint3 GlobalInvocationID: SV_DispatchThreadID,
    uniform BrdfBlock brdfPcb)
{
    float Width, Height;
    brdfPcb.outputTexture.GetDimensions(Width, Height);

    float2 UV = GlobalInvocationID.xy / float2(Width, Height);

    float2 IntegratedBRDF = IntegrateBRDF(UV.x, UV.y);
    brdfPcb.outputTexture[GlobalInvocationID.xy] = IntegratedBRDF;
}