using ILGPU;
using ILGPU.Runtime;
using OpenCvSharp;

namespace Hawkeye.VisionBuilder.Workflow.ColorStudioClasses;

public static class Kernels
{
    static Kernels()
    {
        Init();
        AppDomain.CurrentDomain.ProcessExit +=
            Destruct;
    }

    private static void Destruct(object? sender, EventArgs e)
    {
        _accelerator.Dispose();
        _gpuContext.Dispose();
    }


    private static object _lock = new object();
    private static ILGPU.Context _gpuContext;
    private static Accelerator _accelerator;
    private static Action<Index1D, int, int, ArrayView<byte>, ArrayView<byte>, ArrayView<byte>> _spatialKernel;
    private static Action<Index1D, int, int, ArrayView<byte>, ArrayView<byte>> _dilationKernel;

    public static Action<Index1D, int, int, ArrayView<byte>, ArrayView<byte>, ArrayView<byte>> SpatialKernel => _spatialKernel;

    public static Action<Index1D, int, int, ArrayView<byte>, ArrayView<byte>> Dilation => _dilationKernel;

    public static Accelerator Accelerator
    {
        get => _accelerator;
        set => _accelerator = value;
    }

    static Dictionary<Mat, MemoryBuffer1D<byte,Stride1D.Dense>> _imageCache = new ();

    public static unsafe MemoryBuffer1D<byte, Stride1D.Dense> TransformImageCached(Mat image)
    {
        if (_imageCache.ContainsKey(image))
            return _imageCache[image];
        var imageBytes = new Span<byte>(image.DataPointer, image.Width * image.Height * 3);

        var res=Kernels.Accelerator.Allocate1D<byte>(imageBytes.ToArray());
        _imageCache[image] = res;
        return res;
    }


    static void Init()
    {
        lock (_lock)
        {

            if (_gpuContext != null) return;

            _gpuContext = ILGPU.Context.Create(builder => builder.AllAccelerators());
            _accelerator = _gpuContext.Devices.First(x => x.AcceleratorType == AcceleratorType.Cuda)
                .CreateAccelerator(_gpuContext);
            Console.WriteLine("GPU acceleration initialized");


            _spatialKernel = _accelerator
                .LoadAutoGroupedStreamKernel<Index1D, int, int, ArrayView<byte>, ArrayView<byte>, ArrayView<byte>>(
                    Kernels.TransformSpatialKernel);
            
            _dilationKernel = _accelerator
                .LoadAutoGroupedStreamKernel<Index1D, int, int, ArrayView<byte>,ArrayView<byte>>(
                                       Kernels.DilationKernel);

        }
    }

    private static void TransformSpatialKernel(Index1D index, int w, int mw, ArrayView<byte> candyBytes, ArrayView<byte> maskBytes,  ArrayView<byte> res)
    {
        var i = index;
        var p = i * 3;
        var y = p / w;
        var x = p % w;



        var yp = (int)candyBytes[p + 0];
        var xp = candyBytes[p + 1];
        var vp = (byte)(candyBytes[p + 2] / 5f) * 255;

        var hp = yp + vp;
        var maskColor = maskBytes[hp * mw + xp];
        if (maskColor > 0)
            res[i] = 255;
        else
            res[i] = 0;
    }

    private static void DilationKernel(Index1D index, int w, int size, ArrayView<byte> maskBytes,
        ArrayView<byte> res)
    {
        var i = index;
        var p = i ;
        var y = p / w;
        var x = p % w;

        
        var halfSize = size-1 / 2;

        
        if (maskBytes[p] > 0)
        {
            for (int j = -halfSize; j <= halfSize; j++)
            {
                for (int k = -halfSize; k <= halfSize; k++)
                {
                    var yp = y + j;
                    var xp = x + k;
                    if (yp < 0 || yp >= w || xp < 0 || xp >= w)
                        continue;
                    res[yp * w + xp] = 255;
                }
            }
        }
        
    }

    private static void ErosionKernel(Index1D index, int w, int size, ArrayView<byte> maskBytes,
        ArrayView<byte> res)
    {
        var i = index;
        var p = i ;
        var y = p / w;
        var x = p % w;

        
        var halfSize = size-1 / 2;


        if (maskBytes[p] <255)
        {
            for (int j = -halfSize; j <= halfSize; j++)
            {
                for (int k = -halfSize; k <= halfSize; k++)
                {
                    var yp = y + j;
                    var xp = x + k;
                    if (yp < 0 || yp >= w || xp < 0 || xp >= w)
                        continue;
                    res[yp * w + xp] = 0;
                }
            }
        }
        
    }


}