module Tile.Affine
  ( AffineRankSpace (..),
    Point,
    rowMajor,
    rankOf,
    rankOfMaybe,
    pointOf,
    pointOfMaybe,
    spaceExtent,
    select,
    fixDim,
    ranks,
  )
where

import Control.Monad (guard)
import Tile.Shape (Shape)

data AffineRankSpace = AffineRankSpace
  { AffineRankSpace -> Int
offset :: Int,
    AffineRankSpace -> [Int]
sizes :: [Int],
    AffineRankSpace -> [Int]
strides :: [Int]
  }
  deriving (Int -> AffineRankSpace -> ShowS
[AffineRankSpace] -> ShowS
AffineRankSpace -> String
(Int -> AffineRankSpace -> ShowS)
-> (AffineRankSpace -> String)
-> ([AffineRankSpace] -> ShowS)
-> Show AffineRankSpace
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AffineRankSpace -> ShowS
showsPrec :: Int -> AffineRankSpace -> ShowS
$cshow :: AffineRankSpace -> String
show :: AffineRankSpace -> String
$cshowList :: [AffineRankSpace] -> ShowS
showList :: [AffineRankSpace] -> ShowS
Show, AffineRankSpace -> AffineRankSpace -> Bool
(AffineRankSpace -> AffineRankSpace -> Bool)
-> (AffineRankSpace -> AffineRankSpace -> Bool)
-> Eq AffineRankSpace
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AffineRankSpace -> AffineRankSpace -> Bool
== :: AffineRankSpace -> AffineRankSpace -> Bool
$c/= :: AffineRankSpace -> AffineRankSpace -> Bool
/= :: AffineRankSpace -> AffineRankSpace -> Bool
Eq)

type Point = [Int]

rowMajor :: Shape -> AffineRankSpace
rowMajor :: [Int] -> AffineRankSpace
rowMajor [Int]
shape =
  AffineRankSpace
    { offset :: Int
offset = Int
0,
      sizes :: [Int]
sizes = [Int]
shape,
      strides :: [Int]
strides = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
shape)
    }

rankOf :: AffineRankSpace -> Point -> Int
rankOf :: AffineRankSpace -> [Int] -> Int
rankOf AffineRankSpace
space [Int]
coord =
  case AffineRankSpace -> [Int] -> Maybe Int
rankOfMaybe AffineRankSpace
space [Int]
coord of
    Just Int
rank -> Int
rank
    Maybe Int
Nothing -> String -> Int
forall a. HasCallStack => String -> a
error String
"rankOf: coordinate dimension mismatch"

rankOfMaybe :: AffineRankSpace -> Point -> Maybe Int
rankOfMaybe :: AffineRankSpace -> [Int] -> Maybe Int
rankOfMaybe AffineRankSpace
space [Int]
coord
  | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
coord Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (AffineRankSpace -> [Int]
strides AffineRankSpace
space) Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
inBounds [Int]
coord (AffineRankSpace -> [Int]
sizes AffineRankSpace
space)) =
      Int -> Maybe Int
forall a. a -> Maybe a
Just (AffineRankSpace -> Int
offset AffineRankSpace
space Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
coord (AffineRankSpace -> [Int]
strides AffineRankSpace
space)))
  | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing
  where
    inBounds :: a -> a -> Bool
inBounds a
coordinate a
size = a
coordinate a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0 Bool -> Bool -> Bool
&& a
coordinate a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
size

pointOf :: AffineRankSpace -> Int -> Point
pointOf :: AffineRankSpace -> Int -> [Int]
pointOf AffineRankSpace
space Int
rank =
  case AffineRankSpace -> Int -> Maybe [Int]
pointOfMaybe AffineRankSpace
space Int
rank of
    Just [Int]
point -> [Int]
point
    Maybe [Int]
Nothing -> String -> [Int]
forall a. HasCallStack => String -> a
error String
"pointOf: rank outside affine rank space"

pointOfMaybe :: AffineRankSpace -> Int -> Maybe Point
pointOfMaybe :: AffineRankSpace -> Int -> Maybe [Int]
pointOfMaybe AffineRankSpace
space Int
rank
  | Int
rank Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` AffineRankSpace -> [Int]
ranks AffineRankSpace
space =
      [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
coordinate (AffineRankSpace -> [Int]
strides AffineRankSpace
space) (AffineRankSpace -> [Int]
sizes AffineRankSpace
space))
  | Bool
otherwise = Maybe [Int]
forall a. Maybe a
Nothing
  where
    relativeRank :: Int
relativeRank = Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- AffineRankSpace -> Int
offset AffineRankSpace
space

    coordinate :: Int -> Int -> Int
coordinate Int
stride Int
size =
      (Int
relativeRank Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
stride) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
size

spaceExtent :: AffineRankSpace -> Int
spaceExtent :: AffineRankSpace -> Int
spaceExtent AffineRankSpace
space =
  [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (AffineRankSpace -> [Int]
sizes AffineRankSpace
space)

select :: AffineRankSpace -> Int -> Int -> Int -> Int -> Maybe AffineRankSpace
select :: AffineRankSpace
-> Int -> Int -> Int -> Int -> Maybe AffineRankSpace
select AffineRankSpace
space Int
dim Int
begin Int
end Int
step = do
  let shp :: [Int]
shp = AffineRankSpace -> [Int]
sizes AffineRankSpace
space
      sts :: [Int]
sts = AffineRankSpace -> [Int]
strides AffineRankSpace
space

  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
dim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0)
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
dim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shp)
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
step Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)

  let extent :: Int
extent = [Int]
shp [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
dim

  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
begin Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0)
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
begin Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
extent)
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
begin)
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
extent)

  let newOffset :: Int
newOffset = AffineRankSpace -> Int
offset AffineRankSpace
space Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
begin Int -> Int -> Int
forall a. Num a => a -> a -> a
* ([Int]
sts [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
dim)
      newSize :: Int
newSize = (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
begin Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
step

  AffineRankSpace -> Maybe AffineRankSpace
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    AffineRankSpace
      { offset :: Int
offset = Int
newOffset,
        sizes :: [Int]
sizes = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
dim Int
newSize [Int]
shp,
        strides :: [Int]
strides = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replace Int
dim ([Int]
sts [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
dim Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
step) [Int]
sts
      }
  where
    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

fixDim :: AffineRankSpace -> Int -> Int -> Maybe AffineRankSpace
fixDim :: AffineRankSpace -> Int -> Int -> Maybe AffineRankSpace
fixDim AffineRankSpace
space Int
dim Int
i = AffineRankSpace
-> Int -> Int -> Int -> Int -> Maybe AffineRankSpace
select AffineRankSpace
space Int
dim Int
i (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
1

ranks :: AffineRankSpace -> [Int]
ranks :: AffineRankSpace -> [Int]
ranks AffineRankSpace
space = ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (AffineRankSpace -> [Int] -> Int
rankOf AffineRankSpace
space) ([Int] -> [[Int]]
forall {a}. (Num a, Enum a) => [a] -> [[a]]
points (AffineRankSpace -> [Int]
sizes AffineRankSpace
space))
  where
    points :: [a] -> [[a]]
points [] = [[]]
    points (a
n : [a]
ns) =
      [ a
i a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rest
      | a
i <- [a
0 .. a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1],
        [a]
rest <- [a] -> [[a]]
points [a]
ns
      ]