using System.Collections.Generic;
using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

public static class PoissonDiscSamplingJobSystem
{
    [BurstCompile]
    public struct PoissonJob : IJob
    {
        public float radius;
        public float2 regionSize;
        public int numSamplesBeforeRejection;
        public Unity.Mathematics.Random rng;

        public NativeList<float2> points;
        public NativeArray<int> grid;
        public int2 gridSize;
        public float cellSize;

        public void Execute()
        {
            NativeList<float2> spawnPoints = new NativeList<float2>(Allocator.Temp);
            float2 center = regionSize / 2f;
            spawnPoints.Add(center);
            points.Add(center);
            int2 cellPos = (int2)(center / cellSize);
            grid[cellPos.y * gridSize.x + cellPos.x] = points.Length;

            while (spawnPoints.Length > 0)
            {
                int spawnIndex = rng.NextInt(spawnPoints.Length);
                float2 spawnCentre = spawnPoints[spawnIndex];
                bool accepted = false;

                for (int i = 0; i < numSamplesBeforeRejection; i++)
                {
                    float angle = rng.NextFloat(0, math.PI * 2);
                    float2 dir = new float2(math.sin(angle), math.cos(angle));
                    float2 candidate = spawnCentre + dir * rng.NextFloat(radius, 2 * radius);

                    if (IsValid(candidate))
                    {
                        points.Add(candidate);
                        spawnPoints.Add(candidate);
                        int2 gridCoord = (int2)(candidate / cellSize);
                        grid[gridCoord.y * gridSize.x + gridCoord.x] = points.Length;
                        accepted = true;
                        break;
                    }
                }

                if (!accepted)
                    spawnPoints.RemoveAtSwapBack(spawnIndex);
            }

            spawnPoints.Dispose();
        }

        bool IsValid(float2 candidate)
        {
            if (candidate.x < 0 || candidate.x >= regionSize.x || candidate.y < 0 || candidate.y >= regionSize.y)
                return false;

            int2 cell = (int2)(candidate / cellSize);
            int2 searchStart = math.max(0, cell - 2);
            int2 searchEnd = math.min(cell + 2, gridSize - 1);

            for (int x = searchStart.x; x <= searchEnd.x; x++)
            {
                for (int y = searchStart.y; y <= searchEnd.y; y++)
                {
                    int index = grid[y * gridSize.x + x] - 1;
                    if (index >= 0)
                    {
                        float2 point = points[index];
                        if (math.distancesq(candidate, point) < radius * radius)
                            return false;
                    }
                }
            }

            return true;
        }
    }

    public static List<Vector2> GeneratePoints(float radius, Vector2 regionSize, int numSamples = 30)
    {
        float cellSize = radius / Mathf.Sqrt(2);
        int2 gridSize = new int2(
            Mathf.CeilToInt(regionSize.x / cellSize),
            Mathf.CeilToInt(regionSize.y / cellSize)
        );

        NativeArray<int> grid = new NativeArray<int>(gridSize.x * gridSize.y, Allocator.TempJob);
        NativeList<float2> points = new NativeList<float2>(Allocator.TempJob);

        var job = new PoissonJob
        {
            radius = radius,
            regionSize = regionSize,
            numSamplesBeforeRejection = numSamples,
            rng = new Unity.Mathematics.Random((uint)UnityEngine.Random.Range(1, int.MaxValue)),
            grid = grid,
            points = points,
            gridSize = gridSize,
            cellSize = cellSize
        };

        job.Run();

        List<Vector2> finalPoints = new List<Vector2>(points.Length);
        for (int i = 0; i < points.Length; i++)
        {
            finalPoints.Add(points[i]);
        }

        grid.Dispose();
        points.Dispose();

        return finalPoints;
    }
}