using System;
using System.Numerics;
using ILGPU;
using ILGPU.Runtime;
using ILGPU.Runtime.CPU;
using ILGPU.Runtime.Cuda;
using ILGPU.Runtime.OpenCL;

namespace SDFMapCreator;

public partial class SdfKernels {
    private Context gpuContext;
    
    public SdfKernels() {
        // Initialize the GPU context
        gpuContext = Context.Create(builder => builder
            .OpenCL()
            .Cuda()
            .CPU()
            .Math(MathMode.Fast32BitOnly)
        );

        Console.WriteLine("Reading available accelerators (CUDA only)...");
        foreach (var device in gpuContext.GetCudaDevices()) {
            using Accelerator accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"{GetInfoString(accelerator)}");
        }
        Console.WriteLine("Reading available accelerators (OpenCL only)...");
        foreach (var device in gpuContext.GetCLDevices()) {
            using Accelerator accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"{GetInfoString(accelerator)}");
        }
        Console.WriteLine("Reading available accelerators (CPU only)...");
        foreach (var device in gpuContext.GetCPUDevices()) {
            using Accelerator accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"{GetInfoString(accelerator)}");
        }
        
    }

    public void SelfMask(Vector3[,] input, out Vector3[,] mask) {
        var dev = gpuContext.GetPreferredDevice(preferCPU:false);
        int width = input.GetLength(0);
        int height = input.GetLength(1);
        mask = new Vector3[width, height];
        using Accelerator accelerator = dev.CreateAccelerator(gpuContext);
        
        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        buffer.CopyFromCPU(input);
        
        using var maskBuffer = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        
        var selfMaskKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D, ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>>(SelfMaskKernel);

        selfMaskKernel(new (width, height), buffer.View, maskBuffer.View);
        
        accelerator.Synchronize();

        mask = maskBuffer.GetAsArray2D();
    }
    
    public void ABMask(Vector3[,] A, Vector3[,] B, out Vector3[,] mask) {
        var dev = gpuContext.GetPreferredDevice(preferCPU:false);
        int width = A.GetLength(0);
        int height = A.GetLength(1);
        mask = new Vector3[width, height];
        using Accelerator accelerator = dev.CreateAccelerator(gpuContext);
        
        using var bufferA = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        bufferA.CopyFromCPU(A);
        
        using var bufferB = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        bufferB.CopyFromCPU(B);
        
        using var maskBuffer = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        
        var abMaskKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>>(ABMaskKernel);

        abMaskKernel(new (width, height), bufferA.View, bufferB.View, maskBuffer.View);
        
        accelerator.Synchronize();

        mask = maskBuffer.GetAsArray2D();
    }
    
    public void EdgeDetect(Vector3[,] mask, out Vector3[,] edge) {
        var dev = gpuContext.GetPreferredDevice(preferCPU:false);
        int width = mask.GetLength(0);
        int height = mask.GetLength(1);
        edge = new Vector3[width, height];
        using Accelerator accelerator = dev.CreateAccelerator(gpuContext);
        
        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        buffer.CopyFromCPU(mask);
        
        var edgeKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>, uint, uint>(EdgeKernel);

        edgeKernel(new (width, height), buffer.View, (uint)width, (uint)height);
        
        accelerator.Synchronize();

        edge = buffer.GetAsArray2D();
    }
    
    public void Sdf(Vector2[] edges, int width, int height, out Vector3[,] sdf) {
        var dev = gpuContext.GetPreferredDevice(preferCPU:false);
        using Accelerator accelerator = dev.CreateAccelerator(gpuContext);
        
        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new (width, height));
        
        using var edgeBuffer = accelerator.Allocate1D<Vector2>(edges.Length);
        edgeBuffer.CopyFromCPU(edges);
        
        var sdfKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView1D<Vector2, Stride1D.Dense>,
            int, int>(SdfKernel);

        sdfKernel(new (width, height), buffer.View, edgeBuffer.View, width, height);
        
        accelerator.Synchronize();

        sdf = buffer.GetAsArray2D();
    }
    
    private static string GetInfoString(Accelerator a)
    {
        StringWriter infoString = new StringWriter();
        a.PrintInformation(infoString);
        return infoString.ToString();
    }
}