content: import System.Random
  import Data.List (sortBy, genericLength, minimumBy)
  import Data.Ord (comparing)
  
  -- A finite list of dimensional accessors tell a KDTree how to get a
  -- Euclidean dimensional value 'b' out of an arbitrary datum 'a'.
  type DimensionalAccessors a b = [a -> b]
  
  -- A binary tree structure of 'a'.
  data Tree a = Node a (Tree a) (Tree a)
              | Empty
  
  instance Show a => Show (Tree a) where
    show Empty = "Empty"
    show (Node value left right) =
      "(" ++ show value ++ " " ++ show left ++ " " ++ show right ++ ")"
  
  -- A k-d tree structure of 'a' with Euclidean dimensions of 'b'.
  data KDTree a b = KDTree (DimensionalAccessors a b) (Tree a)
  
  instance Show a => Show (KDTree a b) where
    show (KDTree _ tree) = "KDTree " ++ show tree
  
  -- The squared Euclidean distance formula.
  sqrDist :: Num b => DimensionalAccessors a b -> a -> a -> b
  sqrDist dims a b = sum $ map square $ zipWith (-) a' b'
    where
      a' = map ($ a) dims
      b' = map ($ b) dims
  
  square :: Num a => a -> a
  square = (^ 2)
  
  -- Insert a value into a k-d tree.
  insert :: Ord b => KDTree a b -> a -> KDTree a b
  insert (KDTree dims tree) value = KDTree dims $ ins (cycle dims) tree
    where
      ins _      Empty                   = Node value Empty Empty
      ins (d:ds) (Node split left right) =
        if d value < d split
        then Node split (ins ds left) right
        else Node split left (ins ds right)
  
  -- Produce an empty k-d tree.
  empty :: DimensionalAccessors a b -> KDTree a b
  empty dims = KDTree dims Empty
  
  -- Produce a k-d tree with one value.
  singleton :: Ord b => DimensionalAccessors a b -> a -> KDTree a b
  singleton dims value = insert (empty dims) value
  
  -- Create a k-d tree from a list of values using the median-finding algorithm.
  fromList :: Ord b => DimensionalAccessors a b -> [a] -> KDTree a b
  fromList dims values = KDTree dims $ fList (cycle dims) values
    where
      fList _      []     = Empty
      fList (d:ds) values =
        let sorted          = sortBy (comparing d) values
            (lower, higher) = splitAt (genericLength sorted `div` 2) sorted
        in case higher of
          []          -> Empty
          median:rest -> Node median (fList ds lower) (fList ds rest)
  
  -- Create a k-d tree from a list of values by repeatedly inserting the values
  -- into a tree. Faster than median-finding, but can create unbalanced trees.
  fromListLinear :: Ord b => DimensionalAccessors a b -> [a] -> KDTree a b
  fromListLinear dims values = foldl insert (empty dims) values
  
  -- Given a k-d tree, find the nearest value to a given value.
  -- Also report how many nodes were visited.
  nearest :: (Ord b, Num b, Integral c) => KDTree a b -> a -> (Maybe a, c)
  nearest (KDTree dims tree) value = near (cycle dims) tree
    where
      dist = sqrDist dims
      -- If we have an empty tree, then return nothing.
      near _      Empty                    = (Nothing, 1)
      -- We hit a leaf node, so it is the current best.
      near _      (Node split Empty Empty) = (Just split, 1)
      near (d:ds) (Node split left right)  =
        -- Move down the tree in the fashion of insertion.
        let dimdist x y    = square (d x - d y)
            splitDist      = dist value split
            hyperPlaneDist = dimdist value split
            bestLeft       = near ds left
            bestRight      = near ds right
            -- maybeThisBest is the node of the side of the split where the value
            -- resides, and maybeOtherBest is the node on the other side of the split.
            ((maybeThisBest, thisCount), (maybeOtherBest, otherCount)) =
              if d value < d split
              then (bestLeft, bestRight)
              else (bestRight, bestLeft)
        in case maybeThisBest of
          Nothing    ->
            -- From the search point (in this case), the hypersphere radius to
            -- the split node is always >= the hyperplane distance, so we will
            -- always check the other side.
            let count = 1 + thisCount + otherCount
            in case maybeOtherBest of
              -- We are currently at a leaf node, so this is the only choice.
              -- It is not strictly necessary to take care of this case
              -- because of the above pattern matching in near.
              Nothing         -> (Just split, count)
              -- We have a node on the other side, so compare
              -- it to the split point to see which is closer.
              Just otherBest ->
                if dist value otherBest < splitDist
                then (maybeOtherBest, count)
                else (Just split, count)
  
          Just thisBest ->
            let thisBestDist = dist value thisBest
                best         =
                  -- Determine which is the closer node of this side.
                  if splitDist < thisBestDist
                  then split
                  else thisBest
                bestDist     = dist value best
            in
              if bestDist < hyperPlaneDist
              -- If the distance to the best node is less than the distance
              -- to the splitting hyperplane, then the current best node is the
              -- only choice.
              then (Just best, 1 + thisCount)
              -- There is a chance that a node on the other side is closer
              -- than the current best.
              else
                let count = 1 + thisCount + otherCount
                in case maybeOtherBest of
                  Nothing        -> (Just best, count)
                  Just otherBest ->
                    if bestDist < dist value otherBest
                    then (Just best, count)
                    else (maybeOtherBest, count)
  
  -- Dimensional accessors for a 2-tuple
  tuple2D :: [(a, a) -> a]
  tuple2D = [fst, snd]
  
  -- Dimensional accessors for a 3-tuple
  tuple3D :: [(a, a, a) -> a]
  tuple3D = [d1, d2, d3]
    where
      d1 (a, _, _) = a
      d2 (_, b, _) = b
      d3 (_, _, c) = c
  
  -- Random 3-tuple generation
  instance (Random a, Random b, Random c) => Random (a, b, c) where
    random gen =
      let (vA, genA) = random gen
          (vB, genB) = random genA
          (vC, genC) = random genB
      in ((vA, vB, vC), genC)
  
    randomR ((lA, lB, lC), (hA, hB, hC)) gen =
      let (vA, genA) = randomR (lA, hA) gen
          (vB, genB) = randomR (lB, hB) genA
          (vC, genC) = randomR (lC, hC) genB
      in ((vA, vB, vC), genC)
  
  printResults :: (Show a, Show b, Show c, Floating c) =>
                  a -> (Maybe a, b) -> DimensionalAccessors a c -> IO ()
  printResults point result dims = do
    let (nearest, visited) = result
    case nearest of
      Nothing    -> putStrLn "Could not find nearest."
      Just value -> do
        let dist = sqrt $ sqrDist dims point value
        putStrLn $ "Point:    " ++ show point
        putStrLn $ "Nearest:  " ++ show value
        putStrLn $ "Distance: " ++ show dist
        putStrLn $ "Visited:  " ++ show visited
        putStrLn ""
  
  -- Naive nearest search used to confirm results.
  linearNearest :: (Ord b, Num b) => DimensionalAccessors a b -> a -> [a] -> Maybe a
  linearNearest _    _     [] = Nothing
  linearNearest dims value xs = Just $ minimumBy (comparing $ sqrDist dims value) xs
  
  main :: IO ()
  main = do
    let wikiValues :: [(Double, Double)]
        wikiValues  = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)]
        wikiTree    = fromList tuple2D wikiValues
        wikiSearch  = (9, 2)
        wikiNearest = nearest wikiTree wikiSearch
    putStrLn "Wikipedia example:"
    printResults wikiSearch wikiNearest tuple2D
  
    let stdGen                = mkStdGen 0
        randRange :: ((Double, Double, Double), (Double, Double, Double))
        randRange             = ((0, 0, 0), (1000, 1000, 1000))
        (randSearch, stdGenB) = randomR randRange stdGen
        randValues            = take 1000 $ randomRs randRange stdGenB
        randTree              = fromList tuple3D randValues
        randNearest           = nearest randTree randSearch
        randNearestLinear     = linearNearest tuple3D randSearch randValues
    putStrLn "1000 random 3D points on the range of [0, 1000):"
    printResults randSearch randNearest tuple3D
    putStrLn "Confirm naive nearest:"
    print randNearestLinear
  