module NearestNeighbour2D (KDTree, empty, fromList, findNearest) where

import Data.List (sort, sortBy, minimumBy)
import Test.QuickCheck
import Data.Maybe (fromJust)

-- | Our dimentions are called fst and snd, seeing as how we keep them in tuples
data Dim = Fst | Snd deriving (Show)

-- | Gets the value of a given point in the given dimention
selector :: Dim -> (((a, a), b) -> a)
selector Fst = fst . fst
selector Snd = snd . fst

-- | Updates a given point so that it's value in the given dimention equals the third argument
updator :: Dim -> (a, a) -> a -> (a, a)
updator Fst = (\(_, y) x -> (x, y))
updator Snd = (\(x, _) y -> (x, y))

-- | Returns a metric between two points
distance :: (Num v) => (v, v) -> (v, v) -> v
distance (x1, y1) (x2, y2) = a*a + b*b where
  a = x2 - x1
  b = y2 - y1

data (Ord v, Num v) => KDTree v e =
  Branch Dim (v, v) e (KDTree v e) (KDTree v e) |
  Empty
  deriving (Show)

-- | An empty tree.
empty :: (Ord v, Num v) => KDTree v e
empty = Empty

-- | Construct a KD tree from a list of points.
fromList :: (Ord v, Num v) => [((v, v), e)] -> KDTree v e
fromList values = make dimlist values where
  dimlist = cycle [Fst, Snd]
  make _ [] = Empty
  make (dim:dimlist) values' =
    Branch dim v e (make dimlist leftvalues) (make dimlist rightvalues) where
    values = sortBy (\x y -> compare ((selector dim) x) ((selector dim) y)) values'
    midindex = length values `div` 2
    (v, e) = values !! midindex
    (leftvalues, rightvalues') = splitAt midindex values
    rightvalues = if null rightvalues' then [] else tail rightvalues'

-- | Find the point in the KD tree which is closest to the given point
findNearest :: (Ord v, Num v, Monad m) => (v, v) -> KDTree v e -> m ((v, v), e)
findNearest _ Empty = fail "Empty tree"
findNearest target branch = return $ sel $ inner target branch where
  sel (x, y, _) = (x, y)
  inner target (Branch dim v e left right) =
    if radiusIntersectsFarSide
      then (if fardist < neardist then (farv, fare, fardist) else (nearv, neare, neardist))
      else (nearv, neare, neardist) where
    (farv, fare, fardist) = best far
    (nearv, neare, neardist) = if mydist < neardist' then (v, e, mydist) else (nearv', neare', neardist')
    (nearv', neare', neardist') = best near
    mydist = distance v target

    best side =
      case side of
        Empty -> (v, e, mydist)
        otherwise -> inner target side

    radiusIntersectsFarSide =
      distance closestfarpoint target < neardist
    closestfarpoint = (updator dim) target $ selector dim (v, ())

    (near, far) = if (selector dim (target, ())) < (selector dim (v, ())) then (left, right) else (right, left)

{-
 - Testing...

-- | Finds the closest point in a list of points, returning that point and the
-- distance from the target
simpleNN :: (Ord v, Num v) => [(v, v)] -> (v, v) -> (v, v)
simpleNN [] target = error "Empty list to simpleNN"
simpleNN values target =
  fst $ minimumBy (\x y -> compare (snd x) (snd y)) $ map (\x -> (x, distance target x)) values

prop_CorrectFloat :: [(Float, Float)] -> (Float, Float) -> Property
prop_CorrectFloat xs v =
  not (null xs) ==>
  distance v a == distance v b where
  a = simpleNN xs v
  b = (fst $ fromJust $ findNearest v $ fromList $ map (\x -> (x, ())) xs)

prop_CorrectInt :: [(Int, Int)] -> (Int, Int) -> Property
prop_CorrectInt xs v =
  not (null xs) ==>
  distance v a == distance v b where
  a = simpleNN xs v
  b = (fst $ fromJust $ findNearest v $ fromList $ map (\x -> (x, ())) xs)

runTests = do
  check (defaultConfig { configMaxTest = 10000, configMaxFail = 100000 }) prop_CorrectFloat
  check (defaultConfig { configMaxTest = 10000, configMaxFail = 100000 }) prop_CorrectInt

-}