//
//  Shaders.metal
//  CheapRetouch
//
//  Metal shaders for exemplar-based inpainting.
//

#include <metal_stdlib>
using namespace metal;

// MARK: - Common Types

struct VertexOut {
    float4 position [[position]];
    float2 texCoord;
};

// MARK: - Mask Operations

// Dilate mask by expanding white regions
kernel void dilateMask(
    texture2d<float, access::read> inMask [[texture(0)]],
    texture2d<float, access::write> outMask [[texture(1)]],
    constant int &radius [[buffer(0)]],
    uint2 gid [[thread_position_in_grid]]
) {
    if (gid.x >= outMask.get_width() || gid.y >= outMask.get_height()) {
        return;
    }

    float maxValue = 0.0;

    for (int dy = -radius; dy <= radius; dy++) {
        for (int dx = -radius; dx <= radius; dx++) {
            int2 samplePos = int2(gid) + int2(dx, dy);

            if (samplePos.x >= 0 && samplePos.x < int(inMask.get_width()) &&
                samplePos.y >= 0 && samplePos.y < int(inMask.get_height())) {
                float value = inMask.read(uint2(samplePos)).r;
                maxValue = max(maxValue, value);
            }
        }
    }

    outMask.write(float4(maxValue, maxValue, maxValue, 1.0), gid);
}

// Gaussian blur for mask feathering
kernel void gaussianBlur(
    texture2d<float, access::read> inTexture [[texture(0)]],
    texture2d<float, access::write> outTexture [[texture(1)]],
    constant float *weights [[buffer(0)]],
    constant int &radius [[buffer(1)]],
    constant bool &horizontal [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    if (gid.x >= outTexture.get_width() || gid.y >= outTexture.get_height()) {
        return;
    }

    float4 sum = float4(0.0);
    float weightSum = 0.0;

    for (int i = -radius; i <= radius; i++) {
        int2 offset = horizontal ? int2(i, 0) : int2(0, i);
        int2 samplePos = int2(gid) + offset;

        if (samplePos.x >= 0 && samplePos.x < int(inTexture.get_width()) &&
            samplePos.y >= 0 && samplePos.y < int(inTexture.get_height())) {
            float weight = weights[abs(i)];
            sum += inTexture.read(uint2(samplePos)) * weight;
            weightSum += weight;
        }
    }

    outTexture.write(sum / weightSum, gid);
}

// MARK: - Patch Matching

// Compute sum of squared differences between two patches
kernel void computePatchSSD(
    texture2d<float, access::read> sourceTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    texture2d<float, access::write> ssdTexture [[texture(2)]],
    constant int2 &targetPos [[buffer(0)]],
    constant int &patchRadius [[buffer(1)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = sourceTexture.get_width();
    int height = sourceTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    // Skip if this position is in the mask (unknown region)
    float maskValue = maskTexture.read(gid).r;
    if (maskValue > 0.5) {
        ssdTexture.write(float4(1e10), gid);
        return;
    }

    float ssd = 0.0;
    int validPixels = 0;

    for (int dy = -patchRadius; dy <= patchRadius; dy++) {
        for (int dx = -patchRadius; dx <= patchRadius; dx++) {
            int2 sourcePos = int2(gid) + int2(dx, dy);
            int2 targetSamplePos = targetPos + int2(dx, dy);

            // Check bounds
            if (sourcePos.x < 0 || sourcePos.x >= width ||
                sourcePos.y < 0 || sourcePos.y >= height ||
                targetSamplePos.x < 0 || targetSamplePos.x >= width ||
                targetSamplePos.y < 0 || targetSamplePos.y >= height) {
                continue;
            }

            // Only compare known pixels in source
            float sourceMask = maskTexture.read(uint2(sourcePos)).r;
            if (sourceMask > 0.5) {
                continue;
            }

            float4 sourceColor = sourceTexture.read(uint2(sourcePos));
            float4 targetColor = sourceTexture.read(uint2(targetSamplePos));

            float3 diff = sourceColor.rgb - targetColor.rgb;
            ssd += dot(diff, diff);
            validPixels++;
        }
    }

    // Normalize by valid pixel count
    if (validPixels > 0) {
        ssd /= float(validPixels);
    } else {
        ssd = 1e10;
    }

    ssdTexture.write(float4(ssd), gid);
}

// MARK: - Inpainting

// Simple diffusion-based inpainting step
kernel void diffuseInpaint(
    texture2d<float, access::read> inTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    texture2d<float, access::write> outTexture [[texture(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = inTexture.get_width();
    int height = inTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    float maskValue = maskTexture.read(gid).r;

    // If not in mask, copy original
    if (maskValue < 0.5) {
        outTexture.write(inTexture.read(gid), gid);
        return;
    }

    // Diffuse from neighbors
    float4 sum = float4(0.0);
    float weightSum = 0.0;

    int offsets[4][2] = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

    for (int i = 0; i < 4; i++) {
        int2 neighborPos = int2(gid) + int2(offsets[i][0], offsets[i][1]);

        if (neighborPos.x >= 0 && neighborPos.x < width &&
            neighborPos.y >= 0 && neighborPos.y < height) {
            float neighborMask = maskTexture.read(uint2(neighborPos)).r;
            float weight = (neighborMask < 0.5) ? 2.0 : 1.0;
            sum += inTexture.read(uint2(neighborPos)) * weight;
            weightSum += weight;
        }
    }

    if (weightSum > 0.0) {
        outTexture.write(sum / weightSum, gid);
    } else {
        outTexture.write(inTexture.read(gid), gid);
    }
}

// Copy patch from source to target
kernel void copyPatch(
    texture2d<float, access::read> sourceTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    texture2d<float, access::read_write> targetTexture [[texture(2)]],
    constant int2 &sourcePos [[buffer(0)]],
    constant int2 &targetPos [[buffer(1)]],
    constant int &patchRadius [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int patchSize = patchRadius * 2 + 1;

    if (gid.x >= uint(patchSize) || gid.y >= uint(patchSize)) {
        return;
    }

    int2 offset = int2(gid) - int2(patchRadius);
    int2 srcCoord = sourcePos + offset;
    int2 dstCoord = targetPos + offset;

    int width = sourceTexture.get_width();
    int height = sourceTexture.get_height();

    // Check bounds
    if (srcCoord.x < 0 || srcCoord.x >= width ||
        srcCoord.y < 0 || srcCoord.y >= height ||
        dstCoord.x < 0 || dstCoord.x >= width ||
        dstCoord.y < 0 || dstCoord.y >= height) {
        return;
    }

    // Only write to masked (unknown) pixels
    float maskValue = maskTexture.read(uint2(dstCoord)).r;
    if (maskValue > 0.5) {
        float4 sourceColor = sourceTexture.read(uint2(srcCoord));
        targetTexture.write(sourceColor, uint2(dstCoord));
    }
}

// MARK: - Edge-aware blending

kernel void edgeAwareBlend(
    texture2d<float, access::read> inTexture [[texture(0)]],
    texture2d<float, access::read> originalTexture [[texture(1)]],
    texture2d<float, access::read> maskTexture [[texture(2)]],
    texture2d<float, access::write> outTexture [[texture(3)]],
    constant float &blendRadius [[buffer(0)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = inTexture.get_width();
    int height = inTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    float maskValue = maskTexture.read(gid).r;
    float4 inpaintedColor = inTexture.read(gid);
    float4 originalColor = originalTexture.read(gid);

    // If far from mask boundary, use appropriate color directly
    if (maskValue < 0.1) {
        outTexture.write(originalColor, gid);
        return;
    }
    if (maskValue > 0.9) {
        outTexture.write(inpaintedColor, gid);
        return;
    }

    // Blend in the transition zone
    float4 result = mix(originalColor, inpaintedColor, maskValue);
    outTexture.write(result, gid);
}

// MARK: - Uniform Background Fill

// Fill masked region with uniform background color using gradient blending from edges
kernel void uniformFill(
    texture2d<float, access::read> sourceTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    texture2d<float, access::write> outTexture [[texture(2)]],
    constant float4 &avgColor [[buffer(0)]],
    constant float &maxDistance [[buffer(1)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = sourceTexture.get_width();
    int height = sourceTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    float maskValue = maskTexture.read(gid).r;

    // If not in mask, copy original
    if (maskValue < 0.5) {
        outTexture.write(sourceTexture.read(gid), gid);
        return;
    }

    // For masked pixels, fill with average color
    outTexture.write(avgColor, gid);
}

// Sample border pixels and compute average color and variance
kernel void sampleBorder(
    texture2d<float, access::read> sourceTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    device atomic_uint *sumR [[buffer(0)]],
    device atomic_uint *sumG [[buffer(1)]],
    device atomic_uint *sumB [[buffer(2)]],
    device atomic_uint *count [[buffer(3)]],
    device atomic_uint *varSum [[buffer(4)]],
    constant int &borderWidth [[buffer(5)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = sourceTexture.get_width();
    int height = sourceTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    float maskValue = maskTexture.read(gid).r;

    // Only sample pixels just outside the mask (border pixels)
    if (maskValue >= 0.5) {
        return;
    }

    // Check if this pixel is adjacent to a masked pixel
    bool isBorder = false;
    for (int dy = -borderWidth; dy <= borderWidth && !isBorder; dy++) {
        for (int dx = -borderWidth; dx <= borderWidth && !isBorder; dx++) {
            int2 neighborPos = int2(gid) + int2(dx, dy);
            if (neighborPos.x >= 0 && neighborPos.x < width &&
                neighborPos.y >= 0 && neighborPos.y < height) {
                if (maskTexture.read(uint2(neighborPos)).r >= 0.5) {
                    isBorder = true;
                }
            }
        }
    }

    if (!isBorder) {
        return;
    }

    float4 color = sourceTexture.read(gid);
    
    // Accumulate color values (scaled to avoid precision issues)
    uint r = uint(color.r * 1000.0);
    uint g = uint(color.g * 1000.0);
    uint b = uint(color.b * 1000.0);
    
    atomic_fetch_add_explicit(sumR, r, memory_order_relaxed);
    atomic_fetch_add_explicit(sumG, g, memory_order_relaxed);
    atomic_fetch_add_explicit(sumB, b, memory_order_relaxed);
    atomic_fetch_add_explicit(count, 1, memory_order_relaxed);
}

// Gradient fill from edges - samples from multiple directions with smooth blending
kernel void gradientFill(
    texture2d<float, access::read> sourceTexture [[texture(0)]],
    texture2d<float, access::read> maskTexture [[texture(1)]],
    texture2d<float, access::read> distanceTexture [[texture(2)]],
    texture2d<float, access::write> outTexture [[texture(3)]],
    constant float &maxDist [[buffer(0)]],
    uint2 gid [[thread_position_in_grid]]
) {
    int width = sourceTexture.get_width();
    int height = sourceTexture.get_height();

    if (gid.x >= uint(width) || gid.y >= uint(height)) {
        return;
    }

    float maskValue = maskTexture.read(gid).r;

    // If not in mask, copy original
    if (maskValue < 0.5) {
        outTexture.write(sourceTexture.read(gid), gid);
        return;
    }

    // Sample colors from 16 directions for smoother blending
    float4 colorSum = float4(0.0);
    float weightSum = 0.0;
    
    // 16 directions: cardinal, diagonal, and intermediate angles
    float2 directions[16] = {
        float2(-1.0, 0.0),   // left
        float2(1.0, 0.0),    // right
        float2(0.0, -1.0),   // up
        float2(0.0, 1.0),    // down
        float2(-1.0, -1.0),  // top-left
        float2(1.0, -1.0),   // top-right
        float2(-1.0, 1.0),   // bottom-left
        float2(1.0, 1.0),    // bottom-right
        float2(-2.0, -1.0),  // intermediate angles
        float2(-1.0, -2.0),
        float2(1.0, -2.0),
        float2(2.0, -1.0),
        float2(2.0, 1.0),
        float2(1.0, 2.0),
        float2(-1.0, 2.0),
        float2(-2.0, 1.0)
    };
    
    for (int d = 0; d < 16; d++) {
        float2 dir = normalize(directions[d]);
        float2 pos = float2(gid);
        float distance = 0.0;
        
        // Walk in this direction until we find a non-masked pixel
        for (int i = 0; i < int(maxDist); i++) {
            pos += dir;
            distance += 1.0;
            
            int2 ipos = int2(pos);
            if (ipos.x < 0 || ipos.x >= width || ipos.y < 0 || ipos.y >= height) {
                break;
            }
            
            float m = maskTexture.read(uint2(ipos)).r;
            if (m < 0.5) {
                // Found edge pixel - sample it with inverse-square distance weight
                // This creates smoother blending than linear
                float4 color = sourceTexture.read(uint2(ipos));
                float weight = 1.0 / ((distance * distance) + 1.0);
                colorSum += color * weight;
                weightSum += weight;
                
                // Also sample a few more pixels in this direction for averaging
                for (int j = 1; j <= 3; j++) {
                    int2 extraPos = ipos + int2(dir * float(j));
                    if (extraPos.x >= 0 && extraPos.x < width && 
                        extraPos.y >= 0 && extraPos.y < height) {
                        float em = maskTexture.read(uint2(extraPos)).r;
                        if (em < 0.5) {
                            float4 extraColor = sourceTexture.read(uint2(extraPos));
                            float extraWeight = 0.5 / ((distance * distance) + float(j * j) + 1.0);
                            colorSum += extraColor * extraWeight;
                            weightSum += extraWeight;
                        }
                    }
                }
                break;
            }
        }
    }
    
    if (weightSum > 0.0) {
        outTexture.write(colorSum / weightSum, gid);
    } else {
        // Fallback - shouldn't happen
        outTexture.write(float4(0.5, 0.5, 0.5, 1.0), gid);
    }
}
