using System.Buffers;
using System.Collections.Concurrent;
using ILGPU;
using ILGPU.Runtime;
using System.Numerics;
using PerfectNumbers.Core;
using System.Runtime.InteropServices;
using System.Runtime.CompilerServices;
using ILGPU.Runtime.OpenCL;
namespace PerfectNumbers.Core.Gpu;

internal enum GpuPow2ModStatus
{
    Success,
    Overflow,
    Unavailable,
}

internal static partial class PrimeOrderGpuHeuristics
{
    private static readonly ConcurrentDictionary<ulong, byte> OverflowedPrimes = new();
    private static readonly ConcurrentDictionary<UInt128, byte> OverflowedPrimesWide = new();
    private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>> Pow2ModKernelCache = new();
    private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>>> Pow2ModKernelWideCache = new();
    private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>> PartialFactorKernelCache = new();
    private static readonly ConcurrentDictionary<Accelerator, OrderKernelLauncher> OrderKernelCache = new();
    private static readonly ConcurrentDictionary<Accelerator, SmallPrimeDeviceCache> SmallPrimeDeviceCaches = new();

    public readonly struct OrderKernelConfig(ulong previousOrder, byte hasPreviousOrder, uint smallFactorLimit, int maxPowChecks, int mode)
    {
        public readonly ulong PreviousOrder = previousOrder;
        public readonly byte HasPreviousOrder = hasPreviousOrder;
        public readonly uint SmallFactorLimit = smallFactorLimit;
        public readonly int MaxPowChecks = maxPowChecks;
        public readonly int Mode = mode;
    }

    public readonly struct OrderKernelBuffers(
        ArrayView1D<ulong, Stride1D.Dense> phiFactors,
        ArrayView1D<int, Stride1D.Dense> phiExponents,
        ArrayView1D<ulong, Stride1D.Dense> workFactors,
        ArrayView1D<int, Stride1D.Dense> workExponents,
        ArrayView1D<ulong, Stride1D.Dense> candidates,
        ArrayView1D<int, Stride1D.Dense> stackIndex,
        ArrayView1D<int, Stride1D.Dense> stackExponent,
        ArrayView1D<ulong, Stride1D.Dense> stackProduct,
        ArrayView1D<ulong, Stride1D.Dense> result,
        ArrayView1D<byte, Stride1D.Dense> status)
    {
        public readonly ArrayView1D<ulong, Stride1D.Dense> PhiFactors = phiFactors;
        public readonly ArrayView1D<int, Stride1D.Dense> PhiExponents = phiExponents;
        public readonly ArrayView1D<ulong, Stride1D.Dense> WorkFactors = workFactors;
        public readonly ArrayView1D<int, Stride1D.Dense> WorkExponents = workExponents;
        public readonly ArrayView1D<ulong, Stride1D.Dense> Candidates = candidates;
        public readonly ArrayView1D<int, Stride1D.Dense> StackIndex = stackIndex;
        public readonly ArrayView1D<int, Stride1D.Dense> StackExponent = stackExponent;
        public readonly ArrayView1D<ulong, Stride1D.Dense> StackProduct = stackProduct;
        public readonly ArrayView1D<ulong, Stride1D.Dense> Result = result;
        public readonly ArrayView1D<byte, Stride1D.Dense> Status = status;
    }

    private delegate void OrderKernelLauncher(
        AcceleratorStream stream,
        Index1D extent,
        ulong prime,
        OrderKernelConfig config,
        MontgomeryDivisorData divisor,
        ArrayView1D<uint, Stride1D.Dense> primes,
        ArrayView1D<ulong, Stride1D.Dense> squares,
        int primeCount,
        OrderKernelBuffers buffers);

    private sealed class SmallPrimeDeviceCache
    {
        public MemoryBuffer1D<uint, Stride1D.Dense>? Primes;
        public MemoryBuffer1D<ulong, Stride1D.Dense>? Squares;
        public int Count;
    }

    private const int WideStackThreshold = 8;
    private static PrimeOrderGpuCapability s_capability = PrimeOrderGpuCapability.Default;

    private const int Pow2WindowSizeBits = 8;
    private const int Pow2WindowOddPowerCount = 1 << (Pow2WindowSizeBits - 1);
    private const ulong Pow2WindowFallbackThreshold = 32UL;
    private const int HeuristicCandidateLimit = 512;
    private const int HeuristicStackCapacity = 256;

    private const int GpuSmallPrimeFactorSlots = 64;

    internal static ConcurrentDictionary<ulong, byte> OverflowRegistry => OverflowedPrimes;
    internal static ConcurrentDictionary<UInt128, byte> OverflowRegistryWide => OverflowedPrimesWide;

    internal static void OverrideCapabilitiesForTesting(PrimeOrderGpuCapability capability)
    {
        s_capability = capability;
    }

    internal static void ResetCapabilitiesForTesting()
    {
        s_capability = PrimeOrderGpuCapability.Default;
    }

    private static SmallPrimeDeviceCache GetSmallPrimeDeviceCache(Accelerator accelerator)
    {
        return SmallPrimeDeviceCaches.GetOrAdd(accelerator, static acc =>
        {
            uint[] primes = PrimesGenerator.SmallPrimes;
            ulong[] squares = PrimesGenerator.SmallPrimesPow2;
            var primeBuffer = acc.Allocate1D<uint>(primes.Length);
            primeBuffer.View.CopyFromCPU(primes);
            var squareBuffer = acc.Allocate1D<ulong>(squares.Length);
            squareBuffer.View.CopyFromCPU(squares);
            return new SmallPrimeDeviceCache
            {
                Primes = primeBuffer,
                Squares = squareBuffer,
                Count = primes.Length,
            };
        });
    }

    private static Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> GetPartialFactorKernel(Accelerator accelerator)
    {
        return PartialFactorKernelCache.GetOrAdd(accelerator, static accel =>
        {
            var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>(PartialFactorKernel);
            var kernel = KernelUtil.GetKernel(loaded);
            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, int, ulong, uint, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>>();
        });
    }

    public static bool TryPartialFactor(
        ulong value,
        uint limit,
        Span<ulong> primeTargets,
        Span<int> exponentTargets,
        out int factorCount,
        out ulong remaining,
        out bool fullyFactored)
    {
        factorCount = 0;
        remaining = value;
        fullyFactored = false;

        if (primeTargets.Length == 0 || exponentTargets.Length == 0)
        {
            return false;
        }

        primeTargets.Clear();
        exponentTargets.Clear();

        if (!TryLaunchPartialFactorKernel(
                value,
                limit,
                primeTargets,
                exponentTargets,
                out int extracted,
                out ulong leftover,
                out bool kernelFullyFactored))
        {
            return false;
        }

        int capacity = Math.Min(primeTargets.Length, exponentTargets.Length);
        if (extracted > capacity)
        {
            factorCount = 0;
            remaining = value;
            fullyFactored = false;
            return false;
        }

        factorCount = extracted;
        remaining = leftover;
        fullyFactored = kernelFullyFactored && leftover == 1UL;
        return true;
    }

    private static bool TryLaunchPartialFactorKernel(
        ulong value,
        uint limit,
        Span<ulong> primeTargets,
        Span<int> exponentTargets,
        out int factorCount,
        out ulong remaining,
        out bool fullyFactored)
    {
        factorCount = 0;
        remaining = value;
        fullyFactored = false;

        try
        {
            var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
            Accelerator accelerator = lease.Accelerator;
            AcceleratorStream stream = lease.Stream;

            var kernel = GetPartialFactorKernel(accelerator);
            SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator);

            // TODO: We should create / reallocate these buffer only once or if the new required length is bigger than capacity.
            var factorBuffer = accelerator.Allocate1D<ulong>(primeTargets.Length);
            var exponentBuffer = accelerator.Allocate1D<int>(exponentTargets.Length);
            var countBuffer = accelerator.Allocate1D<int>(1);
            var remainingBuffer = accelerator.Allocate1D<ulong>(1);
            var fullyFactoredBuffer = accelerator.Allocate1D<byte>(1);

            // TODO: There is no need to clear these buffers because the kernel will always assign values within the required bounds.
			// factorBuffer.MemSetToZero();
            // exponentBuffer.MemSetToZero();
            // countBuffer.MemSetToZero();
            // remainingBuffer.MemSetToZero();
            // fullyFactoredBuffer.MemSetToZero();

            kernel(
                stream,
                1,
                cache.Primes!.View,
                cache.Squares!.View,
                cache.Count,
                primeTargets.Length,
                value,
                limit,
                factorBuffer.View,
                exponentBuffer.View,
                countBuffer.View,
                remainingBuffer.View,
                fullyFactoredBuffer.View);

            stream.Synchronize();

            countBuffer.View.CopyToCPU(ref factorCount, 1);
            factorCount = Math.Min(factorCount, primeTargets.Length);
            factorBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(primeTargets), primeTargets.Length);
            exponentBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentTargets), exponentTargets.Length);
            remainingBuffer.View.CopyToCPU(ref remaining, 1);

            byte fullyFactoredFlag = 0;
            fullyFactoredBuffer.View.CopyToCPU(ref fullyFactoredFlag, 1);
            fullyFactored = fullyFactoredFlag != 0;

            // TODO: These buffers shouldn't be disposed but rather reused accross call with the assigned accelerator.
            factorBuffer.Dispose();
            exponentBuffer.Dispose();
            countBuffer.Dispose();
            remainingBuffer.Dispose();
            fullyFactoredBuffer.Dispose();
            lease.Dispose();
            return true;
        }
        catch (CLException ex)
        {
            Console.WriteLine($"GPU ERROR ({ex.Error}): {ex.Message}");
            factorCount = 0;
            remaining = value;
            fullyFactored = false;
            return false;
        }
        catch (Exception ex) when (ex is AcceleratorException or InternalCompilerException or NotSupportedException or InvalidOperationException or AggregateException)
        {
            Console.WriteLine($"GPU ERROR: {ex.Message}");
            factorCount = 0;
            remaining = value;
            fullyFactored = false;
            return false;
        }
    }

    public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData)
    {
        Span<ulong> exponents = stackalloc ulong[1];
        Span<ulong> remainders = stackalloc ulong[1];
        exponents[0] = exponent;

        GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders, divisorData);
        remainder = remainders[0];
        return status;
    }

    public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan<ulong> exponents, ulong prime, Span<ulong> remainders, in MontgomeryDivisorData divisorData)
    {
        // This will never occur in production code
        // if (exponents.Length == 0)
        // {
        //     return GpuPow2ModStatus.Success;
        // }

        // This will never occur in production code
        // if (remainders.Length < exponents.Length)
        // {
        //     throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders));
        // }

        // This will never occur in production code
        // if (remainders.Length > exponents.Length)
        //     throw new ArgumentException("Remainder span is longer than the exponent span.", nameof(remainders));

        // We need to clear more, but we save on allocations    
        // Span<ulong> target = remainders[..exponents.Length];
        remainders.Clear();

        // This will never occur in production code
        // if (prime <= 1UL)
        // {
        //     return GpuPow2ModStatus.Unavailable;
        // }

        ConcurrentDictionary<ulong, byte> overflowRegistry = OverflowedPrimes;

        if (overflowRegistry.ContainsKey(prime))
        {
            return GpuPow2ModStatus.Overflow;
        }

        PrimeOrderGpuCapability capability = s_capability;

        if (prime.GetBitLength() > capability.ModulusBits)
        {
            overflowRegistry[prime] = 0;
            return GpuPow2ModStatus.Overflow;
        }

        for (int i = 0; i < exponents.Length; i++)
        {
            if (exponents[i].GetBitLength() > capability.ExponentBits)
            {
                return GpuPow2ModStatus.Overflow;
            }
        }

        bool computed = TryComputeOnGpu(exponents, prime, divisorData, remainders);
        return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable;
    }

    public static GpuPow2ModStatus TryPow2Mod(in UInt128 exponent, in UInt128 prime, out UInt128 remainder)
    {
        Span<UInt128> exponents = stackalloc UInt128[1];
        Span<UInt128> remainders = stackalloc UInt128[1];
        exponents[0] = exponent;

        GpuPow2ModStatus status = TryPow2ModBatch(exponents, prime, remainders);
        remainder = remainders[0];
        return status;
    }

    public static GpuPow2ModStatus TryPow2ModBatch(ReadOnlySpan<UInt128> exponents, UInt128 prime, Span<UInt128> remainders)
    {
        return TryPow2ModBatchInternal(exponents, prime, remainders);
    }

    internal static bool TryCalculateOrder(
        ulong prime,
        ulong? previousOrder,
        PrimeOrderCalculator.PrimeOrderSearchConfig config,
        in MontgomeryDivisorData divisorData,
        out ulong order)
    {
        order = 0UL;

        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
        Accelerator accelerator = lease.Accelerator;
        AcceleratorStream stream = lease.Stream;

        var kernel = GetOrderKernel(accelerator);
            SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator);

		// TODO: These buffers should be created reallocated once and assigned to an accelerator. Callers should use their own
		// thread static cache to prevent other threads from using taking accelerators with pre-allocated buffers if they don't
		// use them.
        var phiFactorBuffer = accelerator.Allocate1D<ulong>(GpuSmallPrimeFactorSlots);
        var phiExponentBuffer = accelerator.Allocate1D<int>(GpuSmallPrimeFactorSlots);
        var workFactorBuffer = accelerator.Allocate1D<ulong>(GpuSmallPrimeFactorSlots);
        var workExponentBuffer = accelerator.Allocate1D<int>(GpuSmallPrimeFactorSlots);
        var candidateBuffer = accelerator.Allocate1D<ulong>(HeuristicCandidateLimit);
        var stackIndexBuffer = accelerator.Allocate1D<int>(HeuristicStackCapacity);
        var stackExponentBuffer = accelerator.Allocate1D<int>(HeuristicStackCapacity);
        var stackProductBuffer = accelerator.Allocate1D<ulong>(HeuristicStackCapacity);
        var resultBuffer = accelerator.Allocate1D<ulong>(1);
        var statusBuffer = accelerator.Allocate1D<byte>(1);

		// TODO: Remove the cleaning after the order kernel is modified to always set the result.
		phiFactorBuffer.MemSetToZero();
        phiExponentBuffer.MemSetToZero();
        workFactorBuffer.MemSetToZero();
        workExponentBuffer.MemSetToZero();
        candidateBuffer.MemSetToZero();
        stackIndexBuffer.MemSetToZero();
        stackExponentBuffer.MemSetToZero();
        stackProductBuffer.MemSetToZero();
        resultBuffer.MemSetToZero();
        statusBuffer.MemSetToZero();

        uint limit = config.SmallFactorLimit == 0 ? uint.MaxValue : config.SmallFactorLimit;
        byte hasPrevious = previousOrder.HasValue ? (byte)1 : (byte)0;
        ulong previousValue = previousOrder ?? 0UL;

        var kernelConfig = new OrderKernelConfig(previousValue, hasPrevious, limit, config.MaxPowChecks, (int)config.Mode);
        var buffers = new OrderKernelBuffers(
            phiFactorBuffer.View,
            phiExponentBuffer.View,
            workFactorBuffer.View,
            workExponentBuffer.View,
            candidateBuffer.View,
            stackIndexBuffer.View,
            stackExponentBuffer.View,
            stackProductBuffer.View,
            resultBuffer.View,
            statusBuffer.View);

        kernel(
            stream,
            1,
            prime,
            kernelConfig,
            divisorData,
            cache.Primes!.View,
            cache.Squares!.View,
            cache.Count,
            buffers);

        stream.Synchronize();

        byte status = 0;
        statusBuffer.View.CopyToCPU(ref status, 1);

        PrimeOrderKernelStatus kernelStatus = (PrimeOrderKernelStatus)status;
        if (kernelStatus == PrimeOrderKernelStatus.Fallback)
        {
            DisposeResources();
            lease.Dispose();
            return false;
        }

        if (kernelStatus == PrimeOrderKernelStatus.PollardOverflow)
        {
            DisposeResources();
            lease.Dispose();
            throw new InvalidOperationException("GPU Pollard Rho stack overflow; increase HeuristicStackCapacity.");
        }

        resultBuffer.View.CopyToCPU(ref order, 1);

        if (kernelStatus == PrimeOrderKernelStatus.FactoringFailure)
        {
            order = 0UL;
        }

        DisposeResources();
        lease.Dispose();

        return order != 0UL;

        void DisposeResources()
        {
            phiFactorBuffer.Dispose();
            phiExponentBuffer.Dispose();
            workFactorBuffer.Dispose();
            workExponentBuffer.Dispose();
            candidateBuffer.Dispose();
            stackIndexBuffer.Dispose();
            stackExponentBuffer.Dispose();
            stackProductBuffer.Dispose();
            resultBuffer.Dispose();
            statusBuffer.Dispose();
        }
    }

    private static OrderKernelLauncher GetOrderKernel(Accelerator accelerator)
    {
        return OrderKernelCache.GetOrAdd(accelerator, static accel =>
        {
            var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, OrderKernelConfig, MontgomeryDivisorData, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, OrderKernelBuffers>(CalculateOrderKernel);
            var kernel = KernelUtil.GetKernel(loaded);
            return kernel.CreateLauncherDelegate<OrderKernelLauncher>();
        });
    }

    private static GpuPow2ModStatus TryPow2ModBatchInternal(ReadOnlySpan<UInt128> exponents, UInt128 prime, Span<UInt128> remainders)
    {
        if (exponents.Length == 0)
        {
            return GpuPow2ModStatus.Success;
        }

        if (remainders.Length < exponents.Length)
        {
            throw new ArgumentException("Remainder span is shorter than the exponent span.", nameof(remainders));
        }

        Span<UInt128> target = remainders.Slice(0, exponents.Length);
        target.Clear();

        ConcurrentDictionary<UInt128, byte> overflowRegistryWide = OverflowedPrimesWide;

        if (overflowRegistryWide.ContainsKey(prime))
        {
            return GpuPow2ModStatus.Overflow;
        }

        PrimeOrderGpuCapability capability = s_capability;

        if (prime.GetBitLength() > capability.ModulusBits)
        {
            overflowRegistryWide[prime] = 0;
            return GpuPow2ModStatus.Overflow;
        }

        for (int i = 0; i < exponents.Length; i++)
        {
            if (exponents[i].GetBitLength() > capability.ExponentBits)
            {
                return GpuPow2ModStatus.Overflow;
            }
        }

        bool computed = TryComputeOnGpuWide(exponents, prime, target);
        return computed ? GpuPow2ModStatus.Success : GpuPow2ModStatus.Unavailable;
    }

    private static bool TryComputeOnGpu(ReadOnlySpan<ulong> exponents, ulong prime, in MontgomeryDivisorData divisorData, Span<ulong> results)
    {
        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
        Accelerator accelerator = lease.Accelerator;
        AcceleratorStream stream = lease.Stream;
        var kernel = GetPow2ModKernel(accelerator);
		// TODO: These buffers should be allocated once per accelerator and only reallocated when the new length exceeds capacity.
		// Modify the callers to use their own pool of buffers per accelerator, so that other threads don't use the accelerators
		// with out preallocated buffers. Share the pool with Pow2ModWide kernel.
        var exponentBuffer = accelerator.Allocate1D<ulong>(exponents.Length);
        var remainderBuffer = accelerator.Allocate1D<ulong>(exponents.Length);

        exponentBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(exponents), exponents.Length);

		// Pow2Mod kernel always assigns the required output elements so we don't need to worry about clearing these.
		// remainderBuffer.MemSetToZero();

		try
		{
			kernel(stream, exponents.Length, exponentBuffer.View, divisorData, remainderBuffer.View);
			stream.Synchronize();
		}
		catch (Exception)
		{
			Console.WriteLine($"Exception for {prime} and exponents: {string.Join(",", exponents.ToArray())}.");
			exponentBuffer.Dispose();
			remainderBuffer.Dispose();
			lease.Dispose();
			throw;
		}

        remainderBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(results), exponents.Length);
        exponentBuffer.Dispose();
        remainderBuffer.Dispose();
        lease.Dispose();
        return true;
    }

    private static void ComputePow2ModCpu(ReadOnlySpan<ulong> exponents, ulong prime, in MontgomeryDivisorData divisorData, Span<ulong> results)
    {
        int length = exponents.Length;
        for (int i = 0; i < length; i++)
        {
            results[i] = Pow2ModCpu(exponents[i], prime, divisorData);
        }
    }

    private static bool TryComputeOnGpuWide(ReadOnlySpan<UInt128> exponents, UInt128 prime, Span<UInt128> results)
    {
        int length = exponents.Length;
        if (length == 0)
        {
            return true;
        }

        GpuUInt128[]? rentedExponents = null;
        GpuUInt128[]? rentedResults = null;
        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);

        try
        {
            Accelerator accelerator = lease.Accelerator;
            AcceleratorStream stream = lease.Stream;
            var kernel = GetPow2ModWideKernel(accelerator);
			// TODO: These buffers should be allocated once per accelerator and only reallocated when the new length exceeds capacity.
			// Modify the callers to use their own pool of buffers per accelerator, so that other threads don't use the accelerators
			// with out preallocated buffers. Share the pool with Pow2Mod kernel.
            var exponentBuffer = accelerator.Allocate1D<GpuUInt128>(length);
            var remainderBuffer = accelerator.Allocate1D<GpuUInt128>(length);

            Span<GpuUInt128> exponentSpan = length <= WideStackThreshold
                ? stackalloc GpuUInt128[length]
                : new Span<GpuUInt128>(rentedExponents = ArrayPool<GpuUInt128>.Shared.Rent(length), 0, length);

            for (int i = 0; i < length; i++)
            {
                exponentSpan[i] = (GpuUInt128)exponents[i];
            }

            exponentBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(exponentSpan), length);
			// Pow2ModWide kernel always assigns the required output elements so we don't need to worry about clearing these.
            // remainderBuffer.MemSetToZero();

            GpuUInt128 modulus = (GpuUInt128)prime;
            kernel(stream, length, exponentBuffer.View, modulus, remainderBuffer.View);

            stream.Synchronize();

            Span<GpuUInt128> resultSpan = length <= WideStackThreshold
                ? stackalloc GpuUInt128[length]
                : new Span<GpuUInt128>(rentedResults = ArrayPool<GpuUInt128>.Shared.Rent(length), 0, length);

            remainderBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(resultSpan), length);

            for (int i = 0; i < length; i++)
            {
                results[i] = (UInt128)resultSpan[i];
            }

            exponentBuffer.Dispose();
            remainderBuffer.Dispose();
            return true;
        }
        catch (Exception)
        {
            Console.WriteLine($"Exception when computing {prime} with exponents: {string.Join(", ", exponents.ToArray())}");
            throw;
        }
        finally
        {
            if (rentedExponents is not null)
            {
                ArrayPool<GpuUInt128>.Shared.Return(rentedExponents, clearArray: false);
            }

            if (rentedResults is not null)
            {
                ArrayPool<GpuUInt128>.Shared.Return(rentedResults, clearArray: false);
            }

            lease.Dispose();
        }
    }

    private static void ComputePow2ModCpuWide(ReadOnlySpan<UInt128> exponents, UInt128 prime, Span<UInt128> results)
    {
        int length = exponents.Length;
        for (int i = 0; i < length; i++)
        {
            results[i] = exponents[i].Pow2MontgomeryModWindowed(prime);
        }
    }

    private static ulong Pow2ModCpu(ulong exponent, ulong modulus, in MontgomeryDivisorData divisorData)
    {
        if (modulus <= 1UL)
        {
            return 0UL;
        }

        return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisorData, exponent);
    }

    private static Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> GetPow2ModKernel(Accelerator accelerator)
    {
        return Pow2ModKernelCache.GetOrAdd(accelerator, static accel =>
        {
            var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernel);
            var kernel = KernelUtil.GetKernel(loaded);
            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
        });
    }




    internal readonly record struct PrimeOrderGpuCapability(int ModulusBits, int ExponentBits)
    {
        public static PrimeOrderGpuCapability Default => new(128, 128);
    }
}