module Data.NumIdr.Array.Array

import Data.List
import Data.List1
import Data.Vect
import Data.Zippable
import Data.NumIdr.PrimArray
import Data.NumIdr.Array.Order
import Data.NumIdr.Array.Coords

%default total


||| Arrays are the central data structure of NumIdr. They are an `n`-dimensional
||| grid of values, where `n` is a value known as the *rank* of the array. Arrays
||| of rank 0 are single values, arrays of rank 1 are vectors, and arrays of rank
||| 2 are matrices.
|||
||| Each array has a *shape*, which is a vector of values giving the dimensions
||| of each axis of the array. The shape is also sometimes used to determine the
||| array's total size.
|||
||| Arrays are indexed by row first, as in the standard mathematical notation for
||| matrices. This is independent of the actual order in which they are stored; the
||| default order is row-major, but this is configurable.
|||
||| @ rk The rank of the array
||| @ s The shape of the array
||| @ a The type of the array's elements
export
data Array : (s : Vect rk Nat) -> (a : Type) -> Type where
  ||| Internally, arrays are stored using Idris's primitive array type. This is
  ||| stored along with the array's shape, and a vector of values called the
  ||| *strides*, which determine how indexes into the internal array should be
  ||| performed. This is how the order of the array is configurable.
  |||
  ||| @ ord The order of the elements of the array
  ||| @ sts The strides of the array
  ||| @ s   The shape of the array
  MkArray : (ord : Order) -> (sts : Vect rk Nat) ->
            (s : Vect rk Nat) -> PrimArray a -> Array s a


--------------------------------------------------------------------------------
-- Properties of arrays
--------------------------------------------------------------------------------


||| Extract the primitive array value.
export
getPrim : Array s a -> PrimArray a
getPrim (MkArray _ _ _ arr) = arr

||| The order of the elements of the array
export
getOrder : Array s a -> Order
getOrder (MkArray ord _ _ _) = ord

||| The strides of the array, returned in the same axis order as in the shape.
export
strides : Array {rk} s a -> Vect rk Nat
strides (MkArray _ sts _ _) = sts

||| The total number of elements of the array
||| This is equivalent to `product s`.
export
size : Array s a -> Nat
size = length . getPrim

||| The shape of the array
export
shape : Array {rk} s a -> Vect rk Nat
shape (MkArray _ _ s _) = s

||| The rank of the array
export
rank : Array s a -> Nat
rank = length . shape



shapeEq : (arr : Array s a) -> s = shape arr
shapeEq (MkArray _ _ _ _) = Refl

-- Get a list of all coordinates
getAllCoords' : Vect rk Nat -> List (Vect rk Nat)
getAllCoords' = traverse (\case Z => []; S n => [0..n])

getAllCoords : (s : Vect rk Nat) -> List (Coords s)
getAllCoords [] = pure []
getAllCoords (Z :: s) = []
getAllCoords (S d :: s) = [| forget (allFins d) :: getAllCoords s |]


--------------------------------------------------------------------------------
-- Array constructors
--------------------------------------------------------------------------------


||| Create an array by repeating a single value.
|||
||| @ s    The shape of the constructed array
||| @ ord  The order of the constructed array
export
repeat' : (s : Vect rk Nat) -> (ord : Order) -> a -> Array s a
repeat' s ord x = MkArray ord (calcStrides ord s) s (constant (product s) x)

||| Create an array by repeating a single value.
||| To specify the order of the array, use `repeat'`.
|||
||| @ s The shape of the constructed array
export
repeat : (s : Vect rk Nat) -> a -> Array s a
repeat s = repeat' s COrder


export
zeros : Num a => (s : Vect rk Nat) -> Array s a
zeros s = repeat s 0

export
ones : Num a => (s : Vect rk Nat) -> Array s a
ones s = repeat s 1

||| Create an array given a vector of its elements. The elements of the vector
||| are arranged into the provided shape using the provided order.
|||
||| @ s   The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromVect' : (s : Vect rk Nat) -> (ord : Order) -> Vect (product s) a -> Array s a
fromVect' s ord v = MkArray ord (calcStrides ord s) s (fromList $ toList v)

||| Create an array given a vector of its elements. The elements of the vector
||| are arranged into the provided shape using row-major order (the last axis is the
||| least significant).
||| To specify the order of the array, use `fromVect'`.
|||
||| @ s The shape of the constructed array
export
fromVect : (s : Vect rk Nat) -> Vect (product s) a -> Array s a
fromVect s = fromVect' s COrder

||| Create an array by taking values from a stream.
|||
||| @ s   The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromStream' : (s : Vect rk Nat) -> (ord : Order) -> Stream a -> Array s a
fromStream' s ord st = MkArray ord (calcStrides ord s) s (fromList $ take (product s) st)

||| Create an array by taking values from a stream.
||| To specify the order of the array, use `fromStream'`.
|||
||| @ s The shape of the constructed array
export
fromStream : (s : Vect rk Nat) -> Stream a -> Array s a
fromStream s = fromStream' s COrder

||| Create an array given a function to generate its elements.
|||
||| @ s   The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromFunction' : (s : Vect rk Nat) -> (ord : Order) -> (Coords s -> a) -> Array s a
fromFunction' s ord f = let sts = calcStrides ord s
                        in  MkArray ord sts s (unsafeFromIns (product s) $
                                     map (\is => (getLocation sts is, f is)) $ getAllCoords s)

||| Create an array given a function to generate its elements.
||| To specify the order of the array, use `fromFunction'`.
|||
||| @ s The shape of the constructed array
export
fromFunction : (s : Vect rk Nat) -> (Coords s -> a) -> Array s a
fromFunction s = fromFunction' s COrder

||| Construct an array using a structure of nested vectors. The elements are arranged
||| to the specified order before being written.
|||
||| @ s   The shape of the constructed array
||| @ ord The order of the constructed array
export
array' : (s : Vect rk Nat) -> (ord : Order) -> Vects s a -> Array s a
array' s ord v = MkArray ord sts s (unsafeFromIns (product s) ins)
  where
    sts : Vect rk Nat
    sts = calcStrides ord s

    ins : List (Nat, a)
    ins = collapse $ mapWithIndex (MkPair . getLocation' sts) v

||| Construct an array using a structure of nested vectors.
||| To explicitly specify the shape and order of the array, use `array'`.
export
array : {s : Vect rk Nat} -> Vects s a -> Array s a
array v = MkArray COrder (calcStrides COrder s) s (fromList $ collapse v)


--------------------------------------------------------------------------------
-- Indexing
--------------------------------------------------------------------------------


||| Index the array using the given `Coords` object.
export
index : Coords s -> Array s a -> a
index is arr = index (getLocation (strides arr) is) (getPrim arr)

||| Index the array using the given `CoordsRange` object.
export
indexRange : (rs : CoordsRange s) -> Array s a -> Array (newShape rs) a
indexRange rs arr = let ord = getOrder arr
                        s = newShape {s = shape arr} $ rewrite sym $ shapeEq arr in rs
                        sts = calcStrides ord s
                        -- TODO: Make this actually typecheck without resorting to this mess
                    in  believe_me $ MkArray ord sts s (unsafeFromIns (product s) $
                                                          map (\(is,is') => (getLocation' sts is', index (getLocation' (strides arr) is) (getPrim arr))) $
                                                          getCoordsList {s = shape arr} $ rewrite sym $ shapeEq arr in rs)


--------------------------------------------------------------------------------
-- Operations on arrays
--------------------------------------------------------------------------------


||| Reshape the array into the given shape and reinterpret it according to
||| the given order.
|||
||| @ s'  The shape to convert the array to
||| @ ord The order to reinterpret the array by
export
reshape' : (s' : Vect rk' Nat) -> (ord : Order) -> Array {rk} s a ->
             product s = product s' => Array s' a
reshape' s' ord' arr = MkArray ord' (calcStrides ord' s') s' (getPrim arr)

||| Reshape the array into the given shape.
|||
||| @ s' The shape to convert the array to
export
reshape : (s' : Vect rk' Nat) -> Array {rk} s a ->
            product s = product s' => Array s' a
reshape s' arr = reshape' s' (getOrder arr) arr


||| Change the internal order of the array's elements.
export
reorder : Order -> Array s a -> Array s a
reorder ord' arr = let s = shape arr
                       sts = calcStrides ord' s
                   in  rewrite shapeEq arr
                   in  MkArray ord' sts _ (unsafeFromIns (product s) $
                              map (\is => (getLocation' sts is, index (getLocation' (strides arr) is) (getPrim arr))) $
                              getAllCoords' s)


export
enumerate' : Array {rk} s a -> List (Vect rk Nat, a)
enumerate' (MkArray _ sts sh p) =
  map (\is => (is, index (getLocation' sts is) p)) (getAllCoords' sh)


export
enumerate : Array {rk} s a -> List (Coords s, a)
enumerate arr = map (\is => (is, index is arr))
                    (rewrite shapeEq arr in getAllCoords $ shape arr)

export
stack : (axis : Fin rk) -> Array {rk} s a -> Array (replaceAt axis d s) a ->
        Array (replaceAt axis (index axis s + d) s) a
stack axis a b = let sA = shape a
                     sB = shape b
                     dA = index axis sA
                     dB = index axis sB
                     s = replaceAt axis (dA + dB) sA
                     sts = calcStrides COrder s
                     ins = map (mapFst $ getLocation' sts . toNats) (enumerate a)
                           ++ map (mapFst $ getLocation' sts . updateAt axis (+dA) . toNats) (enumerate b)
                     -- TODO: prove that the type-level shape and `s` are equivalent
                 in  believe_me $ MkArray COrder sts s (unsafeFromIns (product s) ins)



display : Show a => Array s a -> IO ()
display = printLn . PrimArray.toList . getPrim


--------------------------------------------------------------------------------
-- Implementations
--------------------------------------------------------------------------------


-- Most of these implementations apply the operation pointwise. If there are
-- multiple arrays involved with different orders, then all of the arrays are
-- reordered to match one.


export
Zippable (Array s) where
  zipWith f a b = rewrite shapeEq a
                  in MkArray (getOrder a) (strides a) _ $
                     if getOrder a == getOrder b
                        then unsafeZipWith f (getPrim a) (getPrim b)
                        else unsafeZipWith f (getPrim a) (getPrim $ reorder (getOrder a) b)

  zipWith3 f a b c = rewrite shapeEq a
                     in MkArray (getOrder a) (strides a) _ $
                        if (getOrder a == getOrder b) && (getOrder b == getOrder c)
                           then unsafeZipWith3 f (getPrim a) (getPrim b) (getPrim c)
                           else unsafeZipWith3 f (getPrim a) (getPrim $ reorder (getOrder a) b)
                                                             (getPrim $ reorder (getOrder a) c)

  unzipWith f arr = rewrite shapeEq arr
                    in case unzipWith f (getPrim arr) of
                            (a, b) => (MkArray (getOrder arr) (strides arr) _ a,
                                       MkArray (getOrder arr) (strides arr) _ b)

  unzipWith3 f arr = rewrite shapeEq arr
                     in case unzipWith3 f (getPrim arr) of
                             (a, b, c) =>  (MkArray (getOrder arr) (strides arr) _ a,
                                            MkArray (getOrder arr) (strides arr) _ b,
                                            MkArray (getOrder arr) (strides arr) _ c)

export
Functor (Array s) where
  map f (MkArray ord sts s arr) = MkArray ord sts s (map f arr)

export
{s : _} -> Applicative (Array s) where
  pure = repeat s
  (<*>) = zipWith apply

export
{s : _} -> Monad (Array s) where
  join arr = fromFunction s (\is => index is $ index is arr)


-- Foldable and Traversable operate on the primitive array directly. This means
-- that their operation is dependent on the order of the array.

export
Foldable (Array s) where
  foldl f z = foldl f z . getPrim
  foldr f z = foldr f z . getPrim
  null arr = size arr == Z
  toList = toList . getPrim

export
Traversable (Array s) where
  traverse f (MkArray ord sts s arr) =
    map (MkArray ord sts s) (traverse f arr)


export
Eq a => Eq (Array s a) where
  a == b = if getOrder a == getOrder b
              then unsafeEq (getPrim a) (getPrim b)
              else unsafeEq (getPrim a) (getPrim $ reorder (getOrder a) b)

export
Semigroup a => Semigroup (Array s a) where
  (<+>) = zipWith (<+>)

export
{s : _} -> Monoid a => Monoid (Array s a) where
  neutral = repeat s neutral

-- the shape must be known at runtime due to `fromInteger`. If `fromInteger`
-- were moved into its own interface, this constraint could be removed.

export
{s : _} -> Num a => Num (Array s a) where
  (+) = zipWith (+)
  (*) = zipWith (*)

  fromInteger = repeat s . fromInteger


export
Show a => Show (Array s a) where
  showPrec d arr = let orderedElems = PrimArray.toList $ getPrim $
                         if getOrder arr == COrder then arr else reorder COrder arr
                   in  showCon d "array " $ concat $ insertPunct (shape arr) $ map show orderedElems
    where
      splitWindow : Nat -> List String -> List (List String)
      splitWindow n xs = case splitAt n xs of
                           (xs, []) => [xs]
                           (l1, l2) => l1 :: splitWindow n (assert_smaller xs l2)

      insertPunct : Vect rk Nat -> List String -> List String
      insertPunct [] strs = strs
      insertPunct [d] strs = "[" :: intersperse ", " strs `snoc` "]"
      insertPunct (Z :: s) strs = ["[","]"]
      insertPunct (d :: s) strs =
        let secs = if null strs
                      then List.replicate d ("[]" :: Prelude.Nil)
                      else map (insertPunct s) $ splitWindow (length strs `div` d) strs
        in  "[" :: (concat $ intersperse [", "] secs) `snoc` "]"