-- |
-- Module      : Tile.Neighborhood
-- Description : Grid neighbors in affine rank spaces.
--
-- Neighborhood queries use the logical coordinates of an affine rank
-- space, then map neighboring coordinates back to ranks.
module Tile.Neighborhood
  ( -- * Neighbors
    neighbors,
  )
where

import Tile.Affine
import Tile.Shape

-- | Ranks adjacent to a rank along each logical dimension.
--
-- For each dimension, this returns the valid predecessor and successor
-- ranks when they lie inside the affine rank space.
neighbors :: AffineRankSpace -> Int -> [Int]
neighbors :: AffineRankSpace -> Int -> [Int]
neighbors AffineRankSpace
rankSpace Int
rank =
  let point :: [Int]
point = AffineRankSpace -> Int -> [Int]
pointOf AffineRankSpace
rankSpace Int
rank
   in [ AffineRankSpace -> [Int] -> Int
rankOf AffineRankSpace
rankSpace [Int]
p'
      | [Int]
p' <- [Int] -> [Int] -> [[Int]]
neighborPoints (AffineRankSpace -> [Int]
sizes AffineRankSpace
rankSpace) [Int]
point
      ]

neighborPoints :: Shape -> Point -> [Point]
neighborPoints :: [Int] -> [Int] -> [[Int]]
neighborPoints [Int]
shape [Int]
point =
  [[[Int]]] -> [[Int]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ Int -> [[Int]]
moves Int
dim
    | Int
dim <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]
  where
    moves :: Int -> [[Int]]
moves Int
dim =
      let x :: Int
x = [Int]
point [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
dim
          lo :: [Int]
lo = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
dim (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int]
point
          hi :: [Int]
hi = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
dim (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
point
          extent :: Int
extent = [Int]
shape [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
dim
       in [ [Int]
p
          | (Int
v, [Int]
p) <- [(Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, [Int]
lo), (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, [Int]
hi)],
            Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0,
            Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
extent
          ]

replace :: Int -> a -> [a] -> [a]
replace :: forall a. Int -> a -> [a] -> [a]
replace Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
x] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs