import Data.Maybe (catMaybes)
import qualified Data.Set as Set

(|>) = flip ($)

between :: Ord a => a -> (a, a) -> Bool
between x (a, b) = a <= x && x <= b

class Clamp a where
  clampedBy :: a -> (a, a) -> a

instance Clamp Double where
  clampedBy x (a, b) = max a $ min b x

data Vec2
  = Vec2
  { vec2X :: Double
  , vec2Y :: Double
  }
  deriving (Show, Eq, Ord)

instance Clamp Vec2 where
  clampedBy (Vec2 x y) (Vec2 ax ay, Vec2 bx by) =
    Vec2 (x `clampedBy` (ax, bx)) (y `clampedBy` (ay, by))

distanceSq :: Vec2 -> Vec2 -> Double
distanceSq (Vec2 x y) (Vec2 x' y') =
  let dx = x - x'
      dy = y - y' in
    dx * dx + dy * dy

data AABB
  = AABB
  { aabbXY :: Vec2
  , aabbWH :: Vec2
  }
  deriving (Show, Eq)

aabbSE :: AABB -> Vec2
aabbSE (AABB (Vec2 x y) (Vec2 w h)) =
  Vec2 (x + w) (y + h)

aabbEdges :: AABB -> (Vec2, Vec2)
aabbEdges bounds = (aabbXY bounds, aabbSE bounds)

aabb :: Double -> Double -> Double -> Double -> AABB
aabb x y w h = AABB (Vec2 x y) (Vec2 w h)

quadrisect :: AABB -> (AABB, AABB, AABB, AABB)
quadrisect (AABB (Vec2 x y) (Vec2 w h)) =
  ( AABB (Vec2 (x+0)  (y+0))  (Vec2 w2 h2)
  , AABB (Vec2 (x+w2) (y+0))  (Vec2 w2 h2)
  , AABB (Vec2 (x+0)  (y+h2)) (Vec2 w2 h2)
  , AABB (Vec2 (x+w2) (y+h2)) (Vec2 w2 h2)
  )
  where
    w2 = w / 2
    h2 = h / 2

data Circle
  = Circle
  { circleXY :: Vec2
  , circleR :: Double
  }
  deriving (Show, Eq)

circle :: Double -> Double -> Double -> Circle
circle x y r = Circle (Vec2 x y) r

class Geometric a where
  inBounds :: AABB -> a -> Bool

class Intersect a where
  intersect :: a -> a -> Bool

intersectGet :: Intersect a => a -> a -> Maybe (a, a)
intersectGet x y =
  if intersect x y
    then Just (x, y)
    else Nothing

instance Intersect Circle where
  intersect (Circle p r) (Circle q r') =
    distanceSq p q <= (r + r') ** 2

instance Geometric Vec2 where
  inBounds (AABB (Vec2 x y) (Vec2 w h)) (Vec2 x' y') =
       x' `between` (x, x + w)
    && y' `between` (y, y + h)

instance Geometric Circle where
  inBounds bounds (Circle xy r) =
    distanceSq xy (xy `clampedBy` aabbEdges bounds) <= r ** 2

data QTree a
  = QLeaf AABB [a]
  | QNode AABB (QTree a) (QTree a) (QTree a) (QTree a)
  deriving Show

emptyQT :: AABB -> QTree a
emptyQT bounds = QLeaf bounds []

buildQT :: Geometric a => Int -> Int -> AABB -> [a] -> QTree a
buildQT maxCapacity maxDepth bounds =
  foldr (insert maxCapacity maxDepth) (emptyQT bounds)

nodeBounds :: QTree a -> AABB
nodeBounds (QLeaf bounds _) = bounds
nodeBounds (QNode bounds _ _ _ _) = bounds

subdivide :: Geometric a => QTree a -> QTree a
subdivide (QLeaf bounds items) =
  let (nwb, neb, swb, seb) = quadrisect bounds in
    QNode
      bounds
      (itemsFor nwb)
      (itemsFor neb)
      (itemsFor swb)
      (itemsFor seb)
  where
    itemsFor b =
      QLeaf b $ filter (inBounds b) items
subdivide node = node

insert :: Geometric a => Int -> Int -> a -> QTree a -> QTree a
insert maxCapacity maxDepth item node =
  if inBounds (nodeBounds node) item
    then go 0 node
    else node
  where
    go depth node@(QLeaf bounds items) =
      if length items >= maxCapacity && depth < maxDepth
        then go depth (subdivide node)
        else QLeaf bounds $ item:items
    go depth (QNode bounds nw ne sw se) =
      QNode
        bounds
        (go (depth+1) nw)
        (go (depth+1) ne)
        (go (depth+1) sw)
        (go (depth+1) se)

quadrants :: QTree a -> [[a]]
quadrants = go
  where
    go (QLeaf _ items) = [items]
    go (QNode _ nw ne sw se) =
      go nw ++ go ne ++ go sw ++ go se

data ID a
  = ID Int a
  deriving Show

instance Eq (ID a) where
  ID i _ == ID j _ = i == j

instance Ord (ID a) where
  compare (ID i _) (ID j _) = compare i j

instance Geometric a => Geometric (ID a) where
  inBounds bounds (ID _ x) = inBounds bounds x

instance Intersect a => Intersect (ID a) where
  intersect (ID i x) (ID j y) = i /= j && intersect x y

-- upper triangle of cartesian product; pairwise order is critical down the line
symmetricProductWith :: (a -> a -> b) -> [a] -> [b]
symmetricProductWith f = go
  where
    go [] = []
    go (x:xs) =
      map (f x) xs ++ go xs

enumerate :: [a] -> [ID a]
enumerate = zipWith ID [0..]

unique :: Ord a => [a] -> [a]
unique = Set.toList . Set.fromList

findCollisions :: Intersect a => [a] -> [(a, a)]
findCollisions =
  catMaybes . symmetricProductWith intersectGet

collisionCount :: (Geometric a, Intersect a) => AABB -> [a] -> Int
collisionCount bounds items =
  items
  |> enumerate
  |> buildQT 50 7 bounds
  |> quadrants
  |> concatMap findCollisions
  |> unique
  |> length