-- Operators and types and stuff for extended arrays and spreadsheets,
-- 030103/BjL
-- Under construction...

module Xarray(
       module Ix, module Pord, -- for convencience!
       Array, array, listArray, (!), bounds, indices, elems, assocs, 
       accumArray, (//), accum, ixmap, --defined for standard Array library
       Bounds, meet, join, inBounds, mkarray, finarray, infarray,
       at, when, whenL,
       proj_x, proj_1, proj_xx, proj_1x, proj_x1, proj_12, proj_21,
       column, row, transpose,
       proj_xxx, proj_1xx, proj_x1x, proj_xx1, proj_12x, proj_1x2, proj_x12,
       proj_123, proj_312, proj_231, proj_132, proj_213, proj_321,
       nest_1x, nest_x1, nest_1xx, nest_x1x, nest_xx1,
       nest_12x, nest_1x2, nest_x12,
       foldrows, foldcols, sumrows, sumcols,
       ifA, ifM, cols2matrix, rows2matrix, matrix2rows, matrix2cols,
       foldlA, foldl1A, foldrA, foldr1A,
       foldlM, foldl1M, foldrM, foldr1M,
       sumA, prodA, andA, orA, maxA, minA,
       sumM, prodM, andM, orM, maxM, minM ) where

import qualified List as L
import Maybe
import Ix
import qualified Array as A -- we're building a new, extended array package
                            -- on top of the ordinary Haskell arrays!
                            -- All standard Haskell array
                            -- operations/functions must be prepended by
                            -- "A." below, and we're free to use them
                            -- unqualified for our extended arrays

import Pord

-- A bound can be an array bound or Universe (representing the universal set
-- of appropriate type)

data Bounds a = B (a,a) | Universe deriving Eq

-- Basic operations on array bounds (used to implement data field semantics
-- for array operations. glb, lub (from Pord) are elementwise min and max,
-- respectively):

meet (B (l1,u1)) (B (l2,u2)) = B (lub l1 l2, glb u1 u2)
meet b            Universe   = b
meet Universe     b          = b

join (B (l1,u1)) (B (l2,u2)) = B (glb l1 l2, lub u1 u2)
join _         Universe      = Universe
join Universe _              = Universe

inBounds (B b) i     = inRange b i
inBounds Universe i = True

-- Our extended arrays are either "conventional" Haskell arrays, or infinite
-- arrays whose elements are defined by a function:

data (Ix a) => Array a b = Arr (A.Array a b) | Inf (a -> b) deriving ()

-- We provide versions of all Haskell array operators, plus some more...:

infixl 9  !, //

-- Creating arrays:

-- Datafield style function for creating arrays from an array bound and a
-- function ranging over the index set defined by the bound (essentially
-- same as the hidden constructor MkArray in module Array):

mkarray (B b) f = Arr (A.array b [(i, f i) | i <- range b])
mkarray Universe f = Inf f

-- convenience array creation functions:

finarray b f = mkarray (B b) f
infarray f   = mkarray Universe f

-- Compatibility with ordinary Haskell array creation:

array b l = Arr (A.array b l)

-- List to array (for compatibility with ordinary Haskell arrays):

listArray b vs =  array b (zip (range b) vs)

-- Indexing:

(Arr a)!i = a A.! i
(Inf f)!i = f i

-- Extracting bounds:

bounds (Arr a) = B (A.bounds a)
bounds (Inf f) = Universe

-- for compatibility with ordinary Haskell arrays (extended versions of
-- array library functions):

indices (Arr a) = range (A.bounds a)
indices _       = error "Cannot take indices of an infinite array"

elems (Arr a) = [a A.! i | i <- A.indices a]
elems _       = error "Cannot take elems of an infinite array"

assocs (Arr a) =  [(i, a A.! i) | i <- A.indices a]
assocs _       = error "Cannot take assocs of an infinite array"

(Arr a) // us =  array (A.bounds a)
                       ([(i,a A.! i) | i <- A.indices a L.\\ [i | (i,_) <- us]]
                        ++ us)

(Inf f) // us =  infarray (\i -> let r = lookup i us in
                                   if r == Nothing then f i else fromJust r
                          )

accum f (Arr a) = foldl (\a (i,v) -> a A.// [(i,f (a A.! i) v)]) a
accum f _       = error "Cannot take accum of an infinite array"

accumArray f z b = accum f (array b [(i,z) | i <- range b])

ixmap                 :: (Ix a, Ix b) => (a,a) -> (a -> b) -> Array b c
                                         -> Array a c
ixmap b f (Arr a) = array b [(i, a A.! f i) | i <- range b]
ixmap b f (Inf g) = infarray (g . f)

-- Make our new arrays members of same classes as ordinary Haskell arrays:

instance  (Ix a)          => Functor (Array a) where
  fmap f (Arr a) = (Arr (fmap f a))
  fmap f (Inf g) = (Inf (f . g))

instance  (Ix a, Eq b)  => Eq (Array a b)  where
  (Arr a) == (Arr a') = A.assocs a == A.assocs a'
  (Arr a) == (Inf f)  = False
  (Inf f) == (Arr a)  = False
  _       == _        = error "Cannot compare infinite arrays for equality"

instance  (Ix a, Ord b) => Ord (Array a b)  where
  (Arr a) <= (Arr a') = A.assocs a <= A.assocs a'
  _       <= _        = error "Cannot compare infinite array(s) for inequality"

instance  (Ix a, Show a, Show b) => Show (Array a b)  where
    showsPrec p (Arr a) = showParen (p > 9) (
                          showString "array " .
                          shows (A.bounds a) . showChar ' ' .
                          shows (A.assocs a)                  )
    showsPrec p _       = error "Cannot show infinite array"

instance  (Ix a, Read a, Read b) => Read (Array a b)  where
    readsPrec p = readParen (p > 9)
           (\r -> [(array b as, u) | ("array",s) <- lex r,
                                     (b,t)       <- reads s,
                                     (as,u)      <- reads t   ])

-- Here comes new stuff:

-- We support some special convenience functions for arrays up to dimension
-- three, se further below

-- Some types of extended array restrictions (at, when, whenL):

-- Array restriction is similar to data field restriction, but with
-- some differences necessitated by the fact that we use arrays only.

-- at restricts with (conventional) array bound:

at a b = mkarray ((B b) `meet` (bounds a)) (\i -> a!i)

-- when restricts with general predicate and fills in specified value in
-- "masked" positions, not changing the bounds:

when p a c = mkarray (bounds a) (\i -> if (p i) then a!i else c)

-- whenL restricts with list:
-- (Not the most efficient way to do it, but acceptable as definition):

whenL l a c = when (\i ->i `elem` l) a c

-- Example 1: a `at` ((2,2),(m,n)) selects the 2..m x 2..n-submatrix of the
-- matrix a.

-- Example 2: whenL [(1,1),(1,2),(3,4)] a Nothing returns a matrix with the
-- same extents as a but "blanked out" with Nothing in all cells except
-- (1,1), (1,2), and (3,4).

-- Example 3: when (\i -> a!i /= Just 0) a Nothing returns a matrix with
-- the same extents as a but "blanked out" in all cells i except where 
-- a!i /= Just 0.

-- proj..., projections that select subarrays with lower (or equal)
-- dimensionality (e.g., row, column from matrix). Note that
-- projections can permute indices to create, for instance, transpose-like
-- transformations.

-- Example 1: proj_1x creates a 1D-array from a 2D-array, where dimension 1
-- of the result is dimension 1 of the argument. Dimension 2 in the argument
-- is fixed by a second argument to proj1x. So proj1x selects a column,
-- e.g., "proj1x a 5" selects column 5 from a.

-- Example 2: proj_21 creates a 2-D array from a 2D-array, where dimension 1
-- of the result is dimension 2 of the argument and where dimension 2
-- of the result is dimension 1 of the argument. This is simply the
-- transpose of the argument.

-- Example 3: proj_12 simply returns its argument.

-- Example 4: proj_xx projects out a single array element.

-- Example 5: proj_2x1 creates a 2-D array from a 3D-array, where dimension
-- 1 of the result is dimension 3 from the argument and dimension 2 of the
-- result is dimension 1 of the argument. Dimension 2 in the argument is
-- fixed by a second argument to proj2x1.

-- operations on 1D-arrays for completeness:

proj_x a i = a!i

proj_1 a = a

-- operations on 2D-arrays:

proj_xx a i j = a!(i,j)

-- column

proj_1x (Arr a) i = finarray (l1,u1) (\j -> a A.! (j,i))
                    where ((l1,l2),(u1,u2)) = A.bounds a

proj_1x (Inf f) i = infarray (\j -> f (j,i))

-- row

proj_x1 (Arr a) i = finarray  (l2,u2) (\j -> a A.! (i,j))
                    where ((l1,l2),(u1,u2)) = A.bounds a

proj_x1 (Inf f) i = infarray (\j -> f (i,j))

proj_12 a = a

-- transpose

proj_21 (Arr a)  = finarray ((l2,l1),(u2,u1)) (\(i,j) -> a A.! (j,i))
                   where ((l1,l2),(u1,u2)) = A.bounds a

proj_21 (Inf f)  = infarray (\(i,j) -> f (j,i))

-- some shorthands:

column a i = proj_1x a i

row a i = proj_x1 a i

transpose a = proj_21 a

-- operations on 3D-arrays:

proj_xxx a i j k = a!(i,j,k)

-- rows, columns in different directions in 3-D:

proj_1xx (Arr a) i j = finarray (l1,u1) (\k -> a A.! (k,i,j))
                       where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_1xx (Inf f) i j = infarray (\k -> f (k,i,j))

proj_x1x (Arr a) i j = finarray (l2,u2) (\k -> a A.! (i,k,j))
                       where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_x1x (Inf f) i j = infarray (\k -> f (i,k,j))

proj_xx1 (Arr a) i j = finarray (l3,u3) (\k -> a A.! (i,j,k))
                       where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_xx1 (Inf f) i j = infarray (\k -> f (i,j,k))

-- Submatrices in three different orientations (parallel with coordinate
-- plane):

proj_12x (Arr a) i = finarray ((l1,l2),(u1,u2)) (\(j,k) -> a A.! (j,k,i))
                     where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_12x (Inf f) i = infarray (\(j,k) -> f (j,k,i))

proj_1x2 (Arr a) i = finarray ((l1,l3),(u1,u3)) (\(j,k) -> a A.! (j,i,k))
                     where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_1x2 (Inf f) i = infarray (\(j,k) -> f (j,i,k))

proj_x12 (Arr a) i = finarray ((l2,l3),(u2,u3)) (\(j,k) -> a A.! (i,j,k))
                     where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_x12 (Inf f) i = infarray (\(j,k) -> f (i,j,k))

-- Transpositions of 2D-submatrices can be expressed as composition of
-- submatrix selection + 2D-transpose, thus not included here

-- The six possible 3D-transpositions:

proj_123 a = a

proj_312 (Arr a) = finarray ((l3,l1,l2),(u3,u1,u2)) (\(i,j,k) -> a A.! (k,i,j))
                   where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_312 (Inf f) = infarray (\(i,j,k) -> f (k,i,j))

proj_231 (Arr a) = finarray ((l2,l3,l1),(u2,u3,u1)) (\(i,j,k) -> a A.! (j,k,i))
                   where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_231 (Inf f) = infarray (\(i,j,k) -> f (j,k,i))

proj_132 (Arr a) = finarray ((l1,l3,l2),(u1,u3,u2)) (\(i,j,k) -> a A.! (i,k,j))
                   where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_132 (Inf f) = infarray (\(i,j,k) -> f (i,k,j))

proj_213 (Arr a) = finarray ((l2,l1,l3),(u2,u1,u3)) (\(i,j,k) -> a A.! (j,i,k))
                   where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_213 (Inf f) = infarray (\(i,j,k) -> f (j,i,k))

proj_321 (Arr a) = finarray ((l3,l2,l1),(u3,u2,u1)) (\(i,j,k) -> a A.! (k,j,i))
                   where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

proj_321 (Inf f) = infarray (\(i,j,k) -> f (k,j,i))

-- Elementwise if-then-else on arrays (with selecting predicate over
-- index) - underfined elements are filled with element e:

ifA b a1 a2 e = mkarray ((bounds a1) `join` (bounds a2))
                            (\i -> if b i
                                   then awrap a1 e i
                                   else awrap a2 e i
                            )

-- nest..., functions that convert 2D- and 3D-arrays into nested arrays in
-- various ways. Naming is similar to the proj...-functions. Contrary to the
-- family of proj-functions, though, we do not provide any index-permuting
-- conversions (they can be effectuated by mapping suitable proj-functions
-- over the nested arrays).

-- Examples: nest_1x creates a nested 1D-array of 1D-arrays from a 2D-array,
-- where each element in the new array is a row in the 2D-array. Thus, we
-- have created an array in the "first" (column) dimension, of arrays
-- projected in the remaining dimension. Similarly, nest_x1 creates an array
-- of the columns. nest_1xx takes a 3D-array and creates a 1D-array along
-- the first dimension, whose elements are 2D-"slices" of the 3D-array along
-- dimensions 2 and 3. And so on. There are a total of eight nest-functions:
-- two 2D -> 1D of 1D, three 3D -> 1D of 2D, and three 3D -> 2D of 1D.

-- An intended use of the nest-functions is to map various functions over
-- them: for instance, to get functions that operate over the rows or
-- columns of a matrix. A concrete example is to sum all the rows of a
-- matrix into a 1D-array of sums. Note that some of the proj-functions can
-- be obtained in this way by mapping suitable element-selecting functions.

-- 2D -> 1D of 1D:

nest_1x (Arr a) =
  finarray (l1,u1) (\i -> finarray (l2,u2) (\j -> a A.!(i,j)))
    where ((l1,l2),(u1,u2)) = A.bounds a

nest_1x (Inf f) = infarray (\i -> infarray (\j -> f (i,j)))

nest_x1 (Arr a) =
  finarray (l2,u2) (\i -> finarray (l1,u1) (\j -> a A.!(j,i)))
    where ((l1,l2),(u1,u2)) = A.bounds a

nest_x1 (Inf f) = infarray (\i -> infarray (\j -> f (j,i)))

-- 3D -> 1D of 2D:

nest_1xx (Arr a) =
  finarray (l1,u1) (\i -> finarray ((l2,l3),(u2,u3)) (\(j,k) -> a A.!(i,j,k)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_1xx (Inf f) = infarray (\i -> infarray (\(j,k) -> f (i,j,k)))

nest_x1x (Arr a) =
  finarray (l2,u2) (\i -> finarray ((l1,l3),(u1,u3)) (\(j,k) -> a A.!(j,i,k)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_x1x (Inf f) = infarray (\i -> infarray (\(j,k) -> f (j,i,k)))

nest_xx1 (Arr a) =
  finarray (l3,u3) (\i -> finarray ((l1,l2),(u1,u2)) (\(j,k) -> a A.!(j,k,i)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_xx1 (Inf f) = infarray (\i -> infarray (\(j,k) -> f (j,k,i)))

-- 3D -> 2D of 1D:

nest_12x (Arr a) =
  finarray ((l1,l2),(u1,u2)) (\(i,j) -> finarray (l3,u3) (\k -> a A.!(i,j,k)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_12x (Inf f) = infarray (\(i,j) -> infarray (\k -> f (i,j,k)))

nest_1x2 (Arr a) =
  finarray ((l1,l3),(u1,u3)) (\(i,j) -> finarray (l2,u2) (\k -> a A.!(i,k,j)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_1x2 (Inf f) = infarray (\(i,j) -> infarray (\k -> f (i,k,j)))

nest_x12 (Arr a) =
  finarray ((l2,l3),(u2,u3)) (\(i,j) -> finarray (l1,u1) (\k -> a A.!(k,i,j)))
    where ((l1,l2,l3),(u1,u2,u3)) = A.bounds a

nest_x12 (Inf f) = infarray (\(i,j) -> infarray (\k -> f (k,i,j)))

-- Two useful matrix operations (fold rows/columns) built on top (can serve
-- as templates for how to build similar functions):

foldrows op init a = fmap (foldlA op init) (nest_1x a)

foldcols op init a = fmap (foldlA op init) (nest_x1 a)

-- (foldlA is defined below)

-- Summation of matrix rows/columns is very commonplace, so let's predefine
-- two functions for this:

sumrows a = foldrows (+) 0 a

sumcols a = foldcols (+) 0 a

-- Elementwise if-then-else on Maybe arrays (with Nothing put into undefined
-- positions):

ifM b a1 a2 = ifA b a1 a2 Nothing

-- Help function that "wraps" an array to return element e in case the index
-- is out of bounds:

awrap a e i = if inBounds (bounds a) i then a!i else e

-- Build matrix from 1-D arrays, using them either as rows or columns
-- Requires that arrays have same extents. (Seems most convenient)

cols2matrix []         = error "Cannot build matrix from empty list of vectors"
cols2matrix c@(col:cols) =
  let eqelems [x]               = True
      eqelems (x:y:ys) | x /= y = False
                       | x == y = eqelems (y:ys)
      b = case bounds col of
            Universe -> Universe
            B (l2,u2) -> B ((l2,0),(u2,length c - 1))
  in if eqelems (map bounds c)
     then mkarray b (\(i,j) -> (c!!j)!i)
     else error "Cannot build matrix from vectors with non-conforming bounds"

-- rows2matrix = transpose . cols2matrix -- (Doesn't work, for some reason)
rows2matrix rows = transpose (cols2matrix rows)

-- Split matrix into list of rows and list of columns, respectively

matrix2rows (Arr a) = [row (Arr a) i | i <- range (l1,u1)]
                        where ((l1,l2),(u1,u2)) = (A.bounds a)
matrix2rows (Inf f) = error "Cannot split infinite matrix into list of rows"

matrix2cols (Arr a) = [column (Arr a) i | i <- range (l2,u2)]
                        where ((l1,l2),(u1,u2)) = (A.bounds a)
matrix2cols (Inf f) = error "Cannot split infinite matrix into list of columns"

-- Array folds (use all the elements in the array, works for finite arrays
-- only):

foldlA f e (Arr a) = foldl f e (A.elems a)
foldlA f e _       = error "Cannot apply foldlA to an infinite array"

foldl1A f (Arr a) = foldl1 f (A.elems a)
foldl1A f _       = error "Cannot apply foldl1A to an infinite array"

foldrA f e (Arr a) = foldr f e (A.elems a)
foldrA f e _       = error "Cannot apply foldrA to an infinite array"

foldr1A f (Arr a) = foldr1 f (A.elems a)
foldr1A f _       = error "Cannot apply foldr1A to an infinite array"

-- Folds for arrays of Maybe element type (catMaybes extracts the defined
-- elements from A.elems a, the list of elements in a):

foldlM f e (Arr a) = foldl f e (catMaybes (A.elems a))
foldlM f e _       = error "Cannot apply foldlM to an infinite array"

foldl1M f (Arr a) = foldl1 f (catMaybes (A.elems a))
foldl1M f _       = error "Cannot apply foldl1M to an infinite array"

foldrM f e (Arr a) = foldr f e (catMaybes (A.elems a))
foldrM f e _       = error "Cannot apply foldrM to an infinite array"

foldr1M f (Arr a) = foldr1 f (catMaybes (A.elems a))
foldr1M f _       = error "Cannot apply foldr1M to an infinite array"

-- (What about scans?)

-- Some common specializations:

-- For arrays:

sumA a = foldlA (+) 0 a

prodA a = foldlA (*) 1 a

andA a = foldlA (&&) True a

orA a = foldlA (||) False a

maxA a = foldl1A max a

minA a = foldl1A min a

-- For Maybe arrays:

sumM a = foldlM (+) 0 a

prodM a = foldlM (*) 1 a

andM a = foldlM (&&) True a

orM a = foldlM (||) False a

maxM a = foldl1M max a

minM a = foldl1M min a