using System.Numerics;
using ILGPU;
using ILGPU.Runtime;

namespace SDFMapCreator;

public partial class SdfKernels {
    
    private const float LUMA_THRESHOLD = 0.0f;
    static void SelfMaskKernel(Index2D index, ArrayView2D<Vector3, Stride2D.DenseX> input, ArrayView2D<Vector3, Stride2D.DenseX> mask) {
        var x = index.X;
        var y = index.Y;
        var value = input[x, y];
        var lumaA = value.X;
        var r = lumaA > LUMA_THRESHOLD ? 1f : 0f;
        mask[x, y] = new(r, 0f, 0f);
    }

    static void ABMaskKernel(Index2D index,
        ArrayView2D<Vector3, Stride2D.DenseX> A,
        ArrayView2D<Vector3, Stride2D.DenseX> B,
        ArrayView2D<Vector3, Stride2D.DenseX> mask
    ) {
        var x = index.X;
        var y = index.Y;
        var valueA = A[x, y];
        var valueB = B[x, y];
        var lumaA = valueA.X;
        var lumaB = valueB.X;
        var r = lumaB - lumaA > LUMA_THRESHOLD ? 1f : 0f;
        mask[x, y] = new(r, 0f, 0f);
    }

    static void EdgeKernel(Index2D index,
        ArrayView2D<Vector3, Stride2D.DenseX> mask,
        uint width, uint height
    ) { // early exit if not on mask
        if (mask[index].X == 0f) return;
        var x = index.X;
        var y = index.Y;
        //if we are on the edge of the image, return false
        if (x == 0 || y == 0 || x == width - 1 || y == height - 1) return;

        //check the 3x3 kernel
        for (var xi = x - 1; xi <= x + 1; xi++) {
            for (var yi = y - 1; yi <= y + 1; yi++) {
                if (xi < 0 || xi >= width || yi < 0 || yi >= height)
                    continue; //skip out of bounds pixels
                if (mask[xi, yi].X == 0f)
                    mask[index].Y = 1f; //if we find a black pixel, return true
            }
        }
    }

    static void SdfKernel(Index2D index,
        ArrayView2D<Vector3, Stride2D.DenseX> sdf,
        ArrayView1D<Vector2, Stride1D.Dense> edges,
        int width, int height
    ) {
        Vector2 pos = new((float)index.X / width, (float)index.Y / height);
        var minDist = 2f;
        var count = edges.IntExtent.Size;

        for (var i = 0; i < count; i++) {
            Vector2 edgeNrm = new(edges[i].X / width, edges[i].Y / height);
            var dist = Vector2.Distance(pos, edgeNrm);
            if (dist < minDist) minDist = dist;
        }

        if (minDist > 1f) minDist = 1f;

        sdf[index] = new(minDist);
    }

    static void GradientKernel(Index2D index,
        ArrayView2D<Vector3, Stride2D.DenseX> mask,
        ArrayView2D<Vector3, Stride2D.DenseX> sdfa,
        ArrayView2D<Vector3, Stride2D.DenseX> sdfb,
        ArrayView2D<Vector3, Stride2D.DenseX> gradient
    ) { //early exit if not on mask
        if (mask[index].X == 0f || mask[index].Y > 0f) return;
        var a = sdfa[index].X;
        var b = sdfb[index].X;
        gradient[index] = new(a / (a + b));
    }

    static Vector3 SampleBilinear(ArrayView2D<Vector3, Stride2D.DenseX> image, float x, float y) {
        int width = image.IntExtent.X;
        int height = image.IntExtent.Y;
        
        var x0 = (int)x;
        var y0 = (int)y;
        var x1 = x0 + 1;
        var y1 = y0 + 1;

        if (x0 < 0 || x1 >= width || y0 < 0 || y1 >= height) return Vector3.Zero;

        var a = new Vector2(x - x0, y - y0);
        var b = new Vector2(1f - a.X, 1f - a.Y);
        
        return Vector3.Lerp(
            Vector3.Lerp(image[x0, y0], image[x1, y0], a.X),
            Vector3.Lerp(image[x0, y1], image[x1, y1], a.X),
            a.Y
        );
    }
    
    static void DirectionalBlurKernel(Index2D index,
        ArrayView2D<Vector3, Stride2D.DenseX> image,
        ArrayView2D<Vector3, Stride2D.DenseX> mask,
        ArrayView2D<Vector3, Stride2D.DenseX> output,
        float radius, float step, float sigma,
        int width, int height
    ) {
        var x = index.X;
        var y = index.Y;
        var value = image[x, y];
        var maskValue = mask[x, y];
        if (maskValue.X == 0f) {
            output[x, y] = value;
            return;
        }

        var gradient = Vector2.Zero;
        
        // calculate the gradient
        for (int i = -1; i <= 1; i++) {
            if (x + i < 0 || x + i >= width || y + i < 0 || y + i >= height) continue;
            gradient.X += i * image[x + i, y].X;
            gradient.Y += i * image[x, y + i].X;
        }
        
        /*
        output[x, y] = new Vector3(float.Abs((gradient.X * 0.5f) + 0.5f),
            float.Abs((gradient.Y * 0.5f) + 0.5f), 0.5f);
        return;
        */

        if (gradient == Vector2.Zero) {
            output[x, y] = value;
            return;
        }

        gradient = Vector2.Normalize(gradient);
        float sum = 0;

        // now we follow the direction line and sample the image for length;
        for (var l = -radius; l <= radius; l += step) {
            var xOffset = (gradient.X * l);
            var yOffset = (gradient.Y * l);
            var xSample = x + xOffset;
            var ySample = y + yOffset;

            if (xSample < 0 || xSample >= width || ySample < 0 || ySample >= height) continue;

            var sampleValue = SampleBilinear(image, xSample, ySample);
            var weight = MathF.Exp(-(l * l) / (2f * sigma * sigma));
            output[x, y] += sampleValue * weight;
            sum += weight;
        }

        output[x, y] = output[x, y] / sum;
    }
}