implementation of Rapidly-exploring random tree (algorithm on wikipedia: https://en.wikipedia.org/wiki/Rapidly-exploring_random_tree)

rrt.hs

import qualified Data.List  as L
import Control.Monad (foldM)
import qualified Data.Map as M
import System.Random (randomRIO, RandomGen, randomR, initStdGen )
import Debug.Trace (trace)
import qualified Data.Set as S

type ConfigMap = (Int, Int)
type Vertex = (Int, Int)
type Graph = M.Map Vertex [(Vertex)]

euclideanDistance :: Vertex -> Vertex -> Float
euclideanDistance (x1,y1) (x2,y2) =  (**0.5) . fromIntegral $ (x1-x2)^2 + (y1-y2)^2


addEdge :: Graph -> Vertex -> Vertex -> Graph 
addEdge g a b = foldl (\acc (x,y) -> M.insertWith (\x y -> L.nub (x++y)) x y acc) g [(a,[b]), (b,[])] 
                 
                   
nearestVertex :: Graph -> Vertex -> Vertex
nearestVertex g a = head $ L.sortBy f (M.keys g)
     where f x y = compare (euclideanDistance x a) (euclideanDistance y a) 


initRRT :: Vertex -> Graph 
initRRT qInit = trace ("initRRT " ++ show qInit) $ M.insert qInit [] M.empty 


newConfig :: Vertex -> Vertex -> Int ->  Vertex
newConfig (x_, y_) (x,y) k = let a = atan ((fromIntegral (y - y_)) / (fromIntegral (x - x_)))
                             in trace ("-> " ++ show a) $ (x_ + (round .cos $a)  , y_  + (round . sin $ a))
                                   


sampledConfig :: RandomGen g => g -> ConfigMap -> (Vertex,g)
sampledConfig g (w,h) = let (a,g_) = randomR (0,w) g 
                            (b, g__) = randomR (0,h) g_
                        in ((a,b), g__)

sampledConfigStream :: RandomGen g => g  -> ConfigMap -> [Vertex]
sampledConfigStream g c = L.unfoldr (\g_ -> Just (sampledConfig g_ c)) $ g 

-- TODO put in delta k as a parameter                                
it_ :: Vertex -> Graph -> Graph 
it_ v g = let (x_,y_) = nearestVertex g v
              nk = newConfig (x_,y_) v 2
          in addEdge g (x_,y_) nk 
                     
                     
rrt :: RandomGen g => g -> ConfigMap -> Graph -> Int -> Graph
rrt gen c g k = foldl (\g_ v -> it_ v g_) g $ take k (sampledConfigStream gen c)                 

main = do
         seed <- initStdGen 
         let initrrt = initRRT (10,10)
             updatedrrt = rrt seed (500,500) initrrt 10 
         putStrLn . M.showTree $ updatedrrt