// [新增] 运行时 IBL 烘焙 Shader。
// 使用全屏三角形逐面写入 Cube RTV，不需要额外顶点缓冲。

#pragma pack_matrix(row_major)

static const float IBL_PI = 3.14159265359f;
static const float IBL_EPSILON = 1.0e-5f;

cbuffer CBBake : register(b0) {
    // x=roughness, y=源 Cube 尺寸, z=源 Cube 最大 mip, w=保留
    float4 gBakeFloatParams;
    // x=目标 face, y=采样数, z/w=保留
    uint4 gBakeUIntParams;
};

TextureCube<float4> gSourceEnvironment : register(t0);
SamplerState gSourceSampler : register(s0);

struct VSOutput {
    float4 positionCS : SV_POSITION;
    float2 uv : TEXCOORD0;
};

VSOutput VSFullscreen(uint vertexId : SV_VertexID) {
    VSOutput output;

    // 一个三角形覆盖整个目标视口。
    const float2 positions[3] = {
        float2(-1.0f, -1.0f),
        float2(-1.0f,  3.0f),
        float2( 3.0f, -1.0f)
    };

    const float2 position = positions[vertexId];
    output.positionCS = float4(position, 0.0f, 1.0f);
    output.uv = float2(position.x * 0.5f + 0.5f,
                       0.5f - position.y * 0.5f);
    return output;
}

// [新增] 把目标 Cube 某个面的 UV 映射为世界方向。
// face 顺序必须与 D3D11：+X、-X、+Y、-Y、+Z、-Z 完全一致。
float3 CubeDirection(uint faceIndex, float2 uv) {
    const float2 ndc = float2(uv.x * 2.0f - 1.0f,
                              1.0f - uv.y * 2.0f);

    float3 direction = float3(ndc.x, ndc.y, 1.0f);
    switch (faceIndex) {
        case 0u: direction = float3( 1.0f, ndc.y, -ndc.x); break; // +X
        case 1u: direction = float3(-1.0f, ndc.y,  ndc.x); break; // -X
        case 2u: direction = float3( ndc.x, 1.0f, -ndc.y); break; // +Y
        case 3u: direction = float3( ndc.x,-1.0f,  ndc.y); break; // -Y
        case 4u: direction = float3( ndc.x, ndc.y,  1.0f); break; // +Z
        case 5u: direction = float3(-ndc.x, ndc.y, -1.0f); break; // -Z
        default: break;
    }
    return normalize(direction);
}

float RadicalInverseVdC(uint bits) {
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return float(bits) * 2.3283064365386963e-10f;
}

float2 Hammersley(uint index, uint sampleCount) {
    return float2(float(index) / float(sampleCount), RadicalInverseVdC(index));
}

void BuildBasis(float3 normal, out float3 tangent, out float3 bitangent) {
    const float3 up = abs(normal.z) < 0.999f
        ? float3(0.0f, 0.0f, 1.0f)
        : float3(1.0f, 0.0f, 0.0f);
    tangent = normalize(cross(up, normal));
    bitangent = cross(normal, tangent);
}

float3 CosineSampleHemisphere(float2 xi, float3 normal) {
    const float radius = sqrt(xi.y);
    const float phi = 2.0f * IBL_PI * xi.x;

    float sinePhi;
    float cosinePhi;
    sincos(phi, sinePhi, cosinePhi);

    const float3 localDirection = float3(
        radius * cosinePhi,
        radius * sinePhi,
        sqrt(saturate(1.0f - xi.y)));

    float3 tangent;
    float3 bitangent;
    BuildBasis(normal, tangent, bitangent);
    return normalize(
        tangent * localDirection.x +
        bitangent * localDirection.y +
        normal * localDirection.z);
}

float3 ImportanceSampleGGX(float2 xi, float3 normal, float roughness) {
    const float alpha = max(roughness * roughness, 0.001f);
    const float alpha2 = alpha * alpha;

    const float phi = 2.0f * IBL_PI * xi.x;
    const float cosTheta = sqrt(
        saturate((1.0f - xi.y) /
                 max(1.0f + (alpha2 - 1.0f) * xi.y, IBL_EPSILON)));
    const float sinTheta = sqrt(saturate(1.0f - cosTheta * cosTheta));

    float sinePhi;
    float cosinePhi;
    sincos(phi, sinePhi, cosinePhi);

    const float3 halfVectorTangent = float3(
        cosinePhi * sinTheta,
        sinePhi * sinTheta,
        cosTheta);

    float3 tangent;
    float3 bitangent;
    BuildBasis(normal, tangent, bitangent);
    return normalize(
        tangent * halfVectorTangent.x +
        bitangent * halfVectorTangent.y +
        normal * halfVectorTangent.z);
}

float DistributionGGXForPdf(float nDotH, float roughness) {
    const float alpha = max(roughness * roughness, 0.001f);
    const float alpha2 = alpha * alpha;
    const float denominator = nDotH * nDotH * (alpha2 - 1.0f) + 1.0f;
    return alpha2 / max(IBL_PI * denominator * denominator, IBL_EPSILON);
}

float4 PSIrradiance(VSOutput input) : SV_TARGET {
    const float3 normal = CubeDirection(gBakeUIntParams.x, input.uv);
    const uint sampleCount = max(gBakeUIntParams.y, 1u);

    float3 irradiance = float3(0.0f, 0.0f, 0.0f);
    [loop]
    for (uint sampleIndex = 0u; sampleIndex < sampleCount; ++sampleIndex) {
        const float2 xi = Hammersley(sampleIndex, sampleCount);
        const float3 lightDirection = CosineSampleHemisphere(xi, normal);
        irradiance += gSourceEnvironment.SampleLevel(
            gSourceSampler, lightDirection, 0.0f).rgb;
    }

    // 这里直接保存 Lambert 卷积后的漫反射环境项 E/PI。
    // 因此运行时只需乘 baseColor，不再额外除以 PI。
    irradiance /= float(sampleCount);
    return float4(irradiance, 1.0f);
}

float4 PSPrefilter(VSOutput input) : SV_TARGET {
    const float3 normal = CubeDirection(gBakeUIntParams.x, input.uv);
    const float roughness = saturate(gBakeFloatParams.x);

    // roughness=0 的 mip0 直接保留原环境，避免 GGX 极窄分布的数值误差。
    if (roughness <= 0.001f) {
        return float4(gSourceEnvironment.SampleLevel(
            gSourceSampler, normal, 0.0f).rgb, 1.0f);
    }

    const uint sampleCount = max(gBakeUIntParams.y, 1u);
    const float sourceSize = max(gBakeFloatParams.y, 1.0f);
    const float sourceMaxMip = max(gBakeFloatParams.z, 0.0f);

    const float3 viewDirection = normal;
    float3 prefilteredColor = float3(0.0f, 0.0f, 0.0f);
    float totalWeight = 0.0f;

    const float texelSolidAngle =
        4.0f * IBL_PI / (6.0f * sourceSize * sourceSize);

    [loop]
    for (uint sampleIndex = 0u; sampleIndex < sampleCount; ++sampleIndex) {
        const float2 xi = Hammersley(sampleIndex, sampleCount);
        const float3 halfVector = ImportanceSampleGGX(xi, normal, roughness);
        const float3 lightDirection = normalize(
            2.0f * dot(viewDirection, halfVector) * halfVector - viewDirection);

        const float nDotL = saturate(dot(normal, lightDirection));
        if (nDotL > 0.0f) {
            const float nDotH = saturate(dot(normal, halfVector));
            const float hDotV = saturate(dot(halfVector, viewDirection));
            const float distribution = DistributionGGXForPdf(nDotH, roughness);
            const float pdf = max(
                distribution * nDotH / max(4.0f * hDotV, IBL_EPSILON),
                IBL_EPSILON);

            const float sampleSolidAngle =
                1.0f / (float(sampleCount) * pdf + IBL_EPSILON);
            const float mipLevel = clamp(
                0.5f * log2(max(sampleSolidAngle / texelSolidAngle, IBL_EPSILON)),
                0.0f, sourceMaxMip);

            prefilteredColor += gSourceEnvironment.SampleLevel(
                gSourceSampler, lightDirection, mipLevel).rgb * nDotL;
            totalWeight += nDotL;
        }
    }

    prefilteredColor /= max(totalWeight, IBL_EPSILON);
    return float4(prefilteredColor, 1.0f);
}

float GeometrySchlickGGXIBL(float nDotX, float roughness) {
    const float k = (roughness * roughness) * 0.5f;
    return nDotX / max(nDotX * (1.0f - k) + k, IBL_EPSILON);
}

float GeometrySmithIBL(float nDotV, float nDotL, float roughness) {
    return GeometrySchlickGGXIBL(nDotV, roughness) *
           GeometrySchlickGGXIBL(nDotL, roughness);
}

float2 IntegrateBrdf(float nDotV, float roughness, uint sampleCount) {
    const float3 normal = float3(0.0f, 0.0f, 1.0f);
    const float3 viewDirection = float3(
        sqrt(saturate(1.0f - nDotV * nDotV)), 0.0f, nDotV);

    float scale = 0.0f;
    float bias = 0.0f;

    [loop]
    for (uint sampleIndex = 0u; sampleIndex < sampleCount; ++sampleIndex) {
        const float2 xi = Hammersley(sampleIndex, sampleCount);
        const float3 halfVector = ImportanceSampleGGX(xi, normal, roughness);
        const float3 lightDirection = normalize(
            2.0f * dot(viewDirection, halfVector) * halfVector - viewDirection);

        const float nDotL = saturate(lightDirection.z);
        const float nDotH = saturate(halfVector.z);
        const float vDotH = saturate(dot(viewDirection, halfVector));

        if (nDotL > 0.0f) {
            const float geometry = GeometrySmithIBL(nDotV, nDotL, roughness);
            const float geometryVisibility =
                geometry * vDotH /
                max(nDotH * nDotV, IBL_EPSILON);
            const float fresnel = pow(1.0f - vDotH, 5.0f);

            scale += (1.0f - fresnel) * geometryVisibility;
            bias += fresnel * geometryVisibility;
        }
    }

    return float2(scale, bias) / float(sampleCount);
}

float2 PSBrdfLut(VSOutput input) : SV_TARGET {
    const float nDotV = clamp(input.uv.x, 0.001f, 1.0f);
    const float roughness = clamp(input.uv.y, 0.001f, 1.0f);
    const uint sampleCount = max(gBakeUIntParams.y, 1u);
    return IntegrateBrdf(nDotV, roughness, sampleCount);
}