module MinDistanceSample where

import qualified Control.Monad as M
import qualified Data.Array as A
import qualified Data.Bifunctor as B
import qualified Data.Ix as Ix
import qualified Data.List as L
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import qualified Data.Maybe as My
import qualified Debug.Trace as D
import Diagrams.Prelude
import System.Random.Stateful

k :: Int
k = 10

randomPolarCoord :: (StatefulGen g m, UniformRange n, RealFloat n) => g -> m (V2 n)
randomPolarCoord gen = do
    angle <- uniformRM (0, 1) gen
    distance <- uniformRM (1, 2) gen
    pure $ view (from r2PolarIso) (distance, angle @@ turn)

randomPoints :: RandomGen r => Point V2 Double -> Double -> r -> [Point V2 Double]
randomPoints dims minDistance gen' = runStateGen_ gen' randomPointsM
  where
    randomPointsM gen = do
        startingPoint <- randomPoint
        map (fmap (* bucketSize)) <$> randomPointsRec [startingPoint] (initialGrid startingPoint)
      where
        initialGrid = addPointToGrid M.empty

        gridBounds = (mkP2 0 0, fmap (floor . (/ bucketSize)) dims)

        randomPoint = p2 <$> randomRM ((0, 0), unp2 . fmap fromIntegral . snd $ gridBounds) gen

        bucketSize :: Double
        bucketSize = minDistance / (sqrt 2)

        addPointToGrid grid p = M.insert (floor <$> p) (p) grid

        randomValueFrom xs = (xs NE.!!) <$> uniformRM (0, pred . length $ xs) gen

        randomPointsRec [] grid = pure . M.elems $ grid
        randomPointsRec (x : xs') grid = do
            startingPoint <- randomValueFrom xs
            newPoint <- L.find isValidPoint <$> candidates startingPoint
            randomPointsRec (prependIfJust newPoint $ (if My.isJust newPoint then x : xs' else xsWithout startingPoint)) . gridPlus $ newPoint
          where
            prependIfJust (Just a) xs = a : xs
            prependIfJust Nothing xs = xs

            gridPlus = maybe grid (addPointToGrid grid)

            xs :: NE.NonEmpty (Point V2 Double)
            xs = x NE.:| xs'

            xsWithout x = NE.filter (/= x) xs

            unitVectorsAround =
                V2
                    <$> [-1, 0, 1]
                    <*> [-1, 0, 1]

            candidates startingPoint = do
                candidatePoints <- M.replicateM k . randomPolarCoord $ gen
                pure $ map (startingPoint .-^) candidatePoints

            isValidPoint :: Point V2 Double -> Bool
            isValidPoint p =
                (Ix.inRange gridBounds (floor <$> p))
                    && ( all ((>= 1) . abs . norm . (p .-.))
                            . My.mapMaybe ((grid M.!?) . fmap floor . (p .-^))
                            $ unitVectorsAround
                       )
