-----------------------------------------------------------------------------
-- |
-- Module      :  Finite.Collection
-- Maintainer  :  Felix Klein
--
-- Allows to extend a finite instance from a single bound to a
-- collection of bounds, given as a finite ranged array.
--
-----------------------------------------------------------------------------

{-# LANGUAGE

    MultiParamTypeClasses
  , LambdaCase
  , ImplicitParams

  #-}

-----------------------------------------------------------------------------

module Finite.Collection where

-----------------------------------------------------------------------------

import Finite.Type
  ( T
  , v2t
  , (#<<)
  , FiniteBounds
  )

import Finite.Class
  ( Finite
  , elements
  , offset
  , value
  , index
  )

import Data.Array.IArray
  ( Array
  , Ix
  , (!)
  , inRange
  , assocs
  , range
  , bounds
  )

import Control.Exception
  ( assert
  )

-----------------------------------------------------------------------------

-- | The 'Collection' type provides a set of items, each assigning an
-- index of type @i@ to a value of type @a@.

data Collection i a =
  Item i a
  deriving
    ( -- | Equality can be checked for collections, if the index type
      -- and the elements can be checked for equality.
      Eq
    , -- | Order can be checked for collections, if the index type and
      -- the elements can be oredered.
      Ord
    , -- | Show a collection through its default constructor.
      Show
    )

-----------------------------------------------------------------------------

-- | Collections are used to extend Finite-Type / Context-Bounds pairs
-- to an array of bounds. At the same time the finite type is extended
-- to a collection of items that range over the same set of indices as
-- the bounds. Since the 'FiniteBounds' parameter always gives a
-- finite sized array of bounding parameters, it is guaranteed that
-- the connected collection has a finite bound as well.

instance (Ix i, Finite b a) => Finite (Array i b) (Collection i a) where

  elements :: T (Collection i a) -> Int
elements T (Collection i a)
t =
    [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ ((i, b) -> Int) -> [(i, b)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (T (Collection i a) -> (i, b) -> Int
forall b a i. Finite b a => T (Collection i a) -> (i, b) -> Int
elms T (Collection i a)
t) ([(i, b)] -> [Int]) -> [(i, b)] -> [Int]
forall a b. (a -> b) -> a -> b
$ Array i b -> [(i, b)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs FiniteBounds (Array i b)
Array i b
?bounds

    where
      conv
        :: T (Collection i a) -> T a

      conv :: T (Collection i a) -> T a
conv = T (Collection i a) -> T a
forall a. HasCallStack => a
undefined


      elms
        :: Finite b a => T (Collection i a) -> (i, b) -> Int

      elms :: T (Collection i a) -> (i, b) -> Int
elms T (Collection i a)
t (i
_,b
b) =
        let ?bounds = b
        in T a -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
elements (T a -> Int) -> T a -> Int
forall a b. (a -> b) -> a -> b
$ T (Collection i a) -> T a
forall i a. T (Collection i a) -> T a
conv T (Collection i a)
t

  index :: Collection i a -> Int
index (Item i
j a
v) =
    let
      -- array bounds
      (i
l,i
u) = Array i b -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds FiniteBounds (Array i b)
Array i b
?bounds
      -- list of indicies that appear before j
      ys :: [i]
ys = Bool -> [i] -> [i]
forall a. HasCallStack => Bool -> a -> a
assert ((i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (i
l,i
u) i
j) ([i] -> [i]) -> [i] -> [i]
forall a b. (a -> b) -> a -> b
$ [i] -> [i]
forall a. [a] -> [a]
init ([i] -> [i]) -> [i] -> [i]
forall a b. (a -> b) -> a -> b
$ (i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i
l,i
j)
      -- offset induces by these indices
      o :: Int
o = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (i -> Int) -> [i] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b -> Int
forall b a. Finite b a => a -> b -> Int
elms a
v (b -> Int) -> (i -> b) -> i -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) (FiniteBounds (Array i b)
Array i b
?bounds Array i b -> i -> b
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!)) [i]
ys
      -- index of v with the bounds at position j
      idx :: Int
idx = let ?bounds = ?bounds ! j
            in a -> Int
forall b a. (Finite b a, FiniteBounds b) => a -> Int
index a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- T a -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
offset (T a -> Int) -> a -> Int
forall a b. (T a -> b) -> a -> b
#<< a
v
    in
      Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx

    where
      elms
        :: Finite b a => a -> b -> Int

      elms :: a -> b -> Int
elms a
v b
b =
        let ?bounds = b
        in T a -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
elements (T a -> Int) -> T a -> Int
forall a b. (a -> b) -> a -> b
$ a -> T a
forall a. a -> T a
v2t a
v

  value :: Int -> Collection i a
value Int
n =
    let
      -- elements of the whole collection
      e :: Int
e = T (Collection i a) -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
elements (T (Collection i a) -> Int) -> T (Collection i a) -> Int
forall a b. (a -> b) -> a -> b
$ Collection i a -> T (Collection i a)
forall a. a -> T a
v2t Collection i a
r
      -- array bounds
      b :: (i, i)
b = Array i b -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds FiniteBounds (Array i b)
Array i b
?bounds
      -- target array index and reminder used as sub-index
      (i
j,Int
m) = T a -> Int -> [i] -> (i, Int)
forall i b a.
(Ix i, Finite b a, FiniteBounds (Array i b)) =>
T a -> Int -> [i] -> (i, Int)
position (Collection i a -> T a
forall i a. Collection i a -> T a
conv Collection i a
r) Int
n ((i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i, i)
b)
      -- result
      r :: Collection i a
r = let ?bounds = ?bounds ! j
          in i -> a -> Collection i a
forall i a. i -> a -> Collection i a
Item i
j (a -> Collection i a) -> a -> Collection i a
forall a b. (a -> b) -> a -> b
$ Int -> a
forall b a. (Finite b a, FiniteBounds b) => Int -> a
value (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ T a -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
offset (Collection i a -> T a
forall i a. Collection i a -> T a
conv Collection i a
r))
    in
      Bool -> Collection i a -> Collection i a
forall a. HasCallStack => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
e) Collection i a
r

    where
      conv
        :: Collection i a -> T a

      conv :: Collection i a -> T a
conv = Collection i a -> T a
forall a. HasCallStack => a
undefined


      position
        :: (Ix i, Finite b a, FiniteBounds (Array i b))
        => T a -> Int -> [i] -> (i,Int)

      position :: T a -> Int -> [i] -> (i, Int)
position T a
t Int
n = \case
        []   -> Bool -> (i, Int) -> (i, Int)
forall a. HasCallStack => Bool -> a -> a
assert Bool
False (i, Int)
forall a. HasCallStack => a
undefined
        i
x:[i]
xr ->
          let m :: Int
m = let ?bounds = ?bounds ! x in T a -> Int
forall b a. (Finite b a, FiniteBounds b) => T a -> Int
elements T a
t
          in if Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n then T a -> Int -> [i] -> (i, Int)
forall i b a.
(Ix i, Finite b a, FiniteBounds (Array i b)) =>
T a -> Int -> [i] -> (i, Int)
position T a
t (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
m) [i]
xr else (i
x,Int
n)

-----------------------------------------------------------------------------