#include <metal_stdlib>
using namespace metal;

// Function to get color based on height and color mode
float3 getColorForHeight(float height, int colorMode) {
    float3 finalColor;
    
    switch (colorMode) {
        case 0: {
            // Realistic terrain coloring
            float3 waterColor = float3(0.2, 0.4, 0.8);      // Deep blue for water/lowest areas
            float3 sandColor = float3(0.8, 0.7, 0.5);       // Sandy beige for beaches/low areas
            float3 grassColor = float3(0.3, 0.6, 0.2);      // Green for grasslands
            float3 forestColor = float3(0.2, 0.4, 0.1);     // Dark green for forests
            float3 rockColor = float3(0.5, 0.4, 0.3);       // Brown-gray for rocky areas
            float3 snowColor = float3(0.9, 0.9, 1.0);       // White-blue for snow peaks
            
            // Define elevation thresholds (0.0 to 1.0 range)
            float waterLevel = 0.15;
            float sandLevel = 0.25;
            float grassLevel = 0.45;
            float forestLevel = 0.65;
            float rockLevel = 0.85;
            
            // Smooth transitions between terrain types using mix function
            if (height < waterLevel) {
                finalColor = waterColor;
            } else if (height < sandLevel) {
                float t = (height - waterLevel) / (sandLevel - waterLevel);
                finalColor = mix(waterColor, sandColor, t);
            } else if (height < grassLevel) {
                float t = (height - sandLevel) / (grassLevel - sandLevel);
                finalColor = mix(sandColor, grassColor, t);
            } else if (height < forestLevel) {
                float t = (height - grassLevel) / (forestLevel - grassLevel);
                finalColor = mix(grassColor, forestColor, t);
            } else if (height < rockLevel) {
                float t = (height - forestLevel) / (rockLevel - forestLevel);
                finalColor = mix(forestColor, rockColor, t);
            } else {
                float t = (height - rockLevel) / (1.0 - rockLevel);
                finalColor = mix(rockColor, snowColor, t);
            }
            break;
        }
        
        case 1: {
            // Standard altitude heat map (blue to red gradient)
            float3 coldColor = float3(0.0, 0.0, 1.0);       // Blue for low elevations
            float3 coolColor = float3(0.0, 1.0, 1.0);       // Cyan
            float3 neutralColor = float3(0.0, 1.0, 0.0);    // Green
            float3 warmColor = float3(1.0, 1.0, 0.0);       // Yellow
            float3 hotColor = float3(1.0, 0.0, 0.0);        // Red for high elevations
            
            if (height < 0.25) {
                float t = height / 0.25;
                finalColor = mix(coldColor, coolColor, t);
            } else if (height < 0.5) {
                float t = (height - 0.25) / 0.25;
                finalColor = mix(coolColor, neutralColor, t);
            } else if (height < 0.75) {
                float t = (height - 0.5) / 0.25;
                finalColor = mix(neutralColor, warmColor, t);
            } else {
                float t = (height - 0.75) / 0.25;
                finalColor = mix(warmColor, hotColor, t);
            }
            break;
        }
        
        default: {
            // Fallback to grayscale
            finalColor = float3(height, height, height);
            break;
        }
    }
    
    return finalColor;
}

[[kernel]]
void convertHeightMapToColors(texture2d<float, access::read> imageTexture [[texture(0)]],
                              texture2d<float, access::write> outputTexture [[texture(1)]],
                              constant int& colorMode [[buffer(0)]],
                              constant float& progress [[buffer(1)]],
                              uint2 pixelCoords [[thread_position_in_grid]]) {
    // Skip out-of-bounds threads.
    if (pixelCoords.x >= imageTexture.get_width() || pixelCoords.y >= imageTexture.get_height()) { return; }
    
    // Read the RGB values from the input image texture
    float4 imageData = imageTexture.read(pixelCoords);
    
    // Use the red channel as height (assuming grayscale input)
    float height = dot(imageData.rgb, float3(0.299, 0.587, 0.114)) * 1.5 * progress;
    
    // Get the color for this height using the specified color mode
    float3 finalColor = getColorForHeight(height, colorMode);
    
    // Write the final color to the output texture
    outputTexture.write(float4(finalColor, 1.0), pixelCoords);
}