// TemperatureConductionCS.usf - 温度传导计算着色器

#include "/Engine/Public/Platform.ush"

// -----------------------------------------------------------------------------
// 1. 输入输出缓冲区定义
// -----------------------------------------------------------------------------
// 输入：当前温度缓冲（只读）
Buffer<float> InputTemperatureBuffer;

// 输入：表面体素标志缓冲（只读）
Buffer<uint> SurfaceVoxelFlagsBuffer;

// 输出：新温度缓冲（读写）
RWBuffer<float> OutputTemperatureBuffer;

// 输出：塑性系数缓冲（读写）
RWBuffer<float> OutputPlasticityBuffer;

// -----------------------------------------------------------------------------
// 2. 参数缓冲区（必须与C++ FParameters结构体严格匹配）
// -----------------------------------------------------------------------------
cbuffer Parameters
{
    int3 VoxelResolution;      // 体素网格分辨率
    float ThermalConductivity; // 热传导系数
    float CoolingRate;         // 表面冷却速率
    float AmbientTemperature;  // 环境温度
    float ForgeTemperatureMin; // 锻造温度下限
    float ForgeTemperatureMax; // 锻造温度上限
    float DeltaTime;           // 时间步长
};

// -----------------------------------------------------------------------------
// 3. 工具函数
// -----------------------------------------------------------------------------

// 获取体素索引（3D坐标转1D索引）
int GetVoxelIndex(int3 coord)
{
    return coord.x + coord.y * VoxelResolution.x +
           coord.z * VoxelResolution.x * VoxelResolution.y;
}

// 检查体素坐标是否有效
bool IsValidCoord(int3 coord)
{
    return all(coord >= 0) && all(coord < VoxelResolution);
}

// 根据温度计算塑性系数（Smoothstep函数）
float CalculatePlasticity(float temperature)
{
    float normalized = saturate((temperature - ForgeTemperatureMin) /
                                (ForgeTemperatureMax - ForgeTemperatureMin));
    return normalized * normalized * (3.0 - 2.0 * normalized);
}

// -----------------------------------------------------------------------------
// 4. 计算着色器主函数
// -----------------------------------------------------------------------------
#define THREADGROUP_SIZE_X 8
#define THREADGROUP_SIZE_Y 8
#define THREADGROUP_SIZE_Z 8

[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)]
void MainCompute(
    uint3 GroupID : SV_GroupID,
    uint3 GroupThreadID : SV_GroupThreadID,
    uint3 DispatchThreadID : SV_DispatchThreadID
)
{
    int3 voxelCoord = int3(DispatchThreadID);

    // 边界检查
    if (!IsValidCoord(voxelCoord))
    {
        return;
    }

    int index = GetVoxelIndex(voxelCoord);
    float currentTemp = InputTemperatureBuffer[index];

    // -------------------------------------------------------------------------
    // 5. 热传导计算（6邻域）
    // -------------------------------------------------------------------------
    float neighborTempSum = 0.0;
    int neighborCount = 0;

    // 6个邻居方向
    const int3 neighbors[6] = {
        int3(-1, 0, 0), int3(1, 0, 0),  // X轴
        int3(0, -1, 0), int3(0, 1, 0),  // Y轴
        int3(0, 0, -1), int3(0, 0, 1)   // Z轴
    };

    // 遍历邻居
    for (int i = 0; i < 6; i++)
    {
        int3 neighborCoord = voxelCoord + neighbors[i];
        if (IsValidCoord(neighborCoord))
        {
            int neighborIndex = GetVoxelIndex(neighborCoord);
            neighborTempSum += InputTemperatureBuffer[neighborIndex];
            neighborCount++;
        }
    }

    // 平均邻居温度
    float avgNeighborTemp = neighborTempSum / max(neighborCount, 1);

    // 热传导增量（向平均温度靠拢）
    float heatConductionDelta = ThermalConductivity * (avgNeighborTemp - currentTemp) * DeltaTime;

    // -------------------------------------------------------------------------
    // 6. 表面冷却计算
    // -------------------------------------------------------------------------
    float coolingDelta = 0.0;
    bool isSurface = (SurfaceVoxelFlagsBuffer[index] & 1) != 0;

    if (isSurface)
    {
        // 表面体素向环境温度冷却
        coolingDelta = CoolingRate * (AmbientTemperature - currentTemp) * DeltaTime;
    }

    // -------------------------------------------------------------------------
    // 7. 更新温度和塑性系数
    // -------------------------------------------------------------------------
    float newTemp = currentTemp + heatConductionDelta + coolingDelta;
    OutputTemperatureBuffer[index] = newTemp;

    // 更新塑性系数
    OutputPlasticityBuffer[index] = CalculatePlasticity(newTemp);
}