import System.Random
import Data.List (sortBy, genericLength, minimumBy)
import Data.Ord (comparing)



type DimensionalAccessors a b = [a -> b]


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 ++ ")" 


data KDTree a b = KDTree (DimensionalAccessors a b) (Tree a)

instance Show a => Show (KDTree a b) where
  show (KDTree _ tree) = "KDTree " ++ show tree


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 :: 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)


empty :: DimensionalAccessors a b -> KDTree a b
empty dims = KDTree dims Empty


singleton :: Ord b => DimensionalAccessors a b -> a -> KDTree a b
singleton dims value = insert (empty dims) value


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)



fromListLinear :: Ord b => DimensionalAccessors a b -> [a] -> KDTree a b
fromListLinear dims values = foldl insert (empty dims) values



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
    
    near _      Empty                    = (Nothing, 1)
    
    near _      (Node split Empty Empty) = (Just split, 1)
    near (d:ds) (Node split left right)  =
      
      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, thisCount), (maybeOtherBest, otherCount)) =
            if d value < d split
            then (bestLeft, bestRight)
            else (bestRight, bestLeft)
      in case maybeThisBest of
        Nothing    ->
          
          
          
          let count = 1 + thisCount + otherCount
          in case maybeOtherBest of
            
            
            
            Nothing         -> (Just split, count)
            
            
            Just otherBest ->
              if dist value otherBest < splitDist
              then (maybeOtherBest, count)
              else (Just split, count)

        Just thisBest ->
          let thisBestDist = dist value thisBest
              best         =
                
                if splitDist < thisBestDist
                then split
                else thisBest
              bestDist     = dist value best
          in
            if bestDist < hyperPlaneDist
            
            
            
            then (Just best, 1 + thisCount)
            
            
            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)


tuple2D :: [(a, a) -> a]
tuple2D = [fst, snd]


tuple3D :: [(a, a, a) -> a]
tuple3D = [d1, d2, d3]
  where
    d1 (a, _, _) = a
    d2 (_, b, _) = b
    d3 (_, _, c) = c


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 ""


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