    {-# LANGUAGE CPP #-}

    {- | Sparse matrices.

    We assume the matrices to be very sparse, so we just implement them as
    sorted association lists.

    Most operations are linear in the number of non-zero elements.

    An exception is transposition, which needs to sort the association
    list again; it has the complexity of sorting: @n log n@ where @n@ is
    the number of non-zero elements.

    Another exception is matrix multiplication, of course.

     -}

    module Agda.Termination.SparseMatrix
      ( -- * Basic data types
        Matrix(Matrix)
      , unM
      -- , matrixInvariant  -- Moved to the internal test-suite
      , Size(..)
      , MIx (..)
        -- * Generating and creating matrices
      , fromLists
      , fromIndexList
      , toLists
    --  , Agda.Termination.Matrix.zipWith
      -- , matrix  -- Moved to the internal test-suite
        -- * Combining and querying matrices
      , size
      , square
      , isEmpty
      , isSingleton
      , zipMatrices
      , add
      , intersectWith
      , interAssocWith
      , mul
      , transpose
      , Diagonal(..)
      , toSparseRows
      , supSize
      , zipAssocWith
        -- * Modifying matrices
      , addRow
      , addColumn
      ) where

    import Data.Array
    import Data.Function
    import qualified Data.List as List
    import Data.Maybe
    import Data.Monoid

    import Data.Foldable (Foldable)
    import qualified Data.Foldable as Fold
    import Data.Traversable (Traversable)

    import qualified Text.PrettyPrint.Boxes as Boxes

    import Agda.Termination.Semiring (HasZero(..), Semiring)
    import qualified Agda.Termination.Semiring as Semiring

    import Agda.Utils.Functor
    import Agda.Utils.List
    import Agda.Utils.Maybe
    import Agda.Utils.Monad
    import Agda.Utils.PartialOrd
    import Agda.Utils.Pretty hiding (isEmpty)
    import Agda.Utils.Tuple

    #include "undefined.h"
    import Agda.Utils.Impossible

    ------------------------------------------------------------------------
    -- * Basic data types
    ------------------------------------------------------------------------

    -- | Size of a matrix.

    data Size i = Size
      { rows :: i  -- ^ Number of rows,    @>= 0@.
      , cols :: i  -- ^ Number of columns, @>= 0@.
      }
      deriving (Eq, Ord, Show)

    -- | Type of matrix indices (row, column).

    data MIx i = MIx
      { row :: i  -- ^ Row index,   @1 <= row <= rows@.
      , col :: i  -- ^ Column index @1 <= col <= cols@.
      }
      deriving (Eq, Ord, Show, Ix)

    -- | Convert a 'Size' to a set of bounds suitable for use with
    --   the matrices in this module.

    toBounds :: Num i => Size i -> (MIx i, MIx i)
    toBounds sz = (MIx { row = 1, col = 1 }, MIx { row = rows sz, col = cols sz })

    -- | Type of matrices, parameterised on the type of values.
    --
    --   Sparse matrices are implemented as an ordered association list,
    --   mapping coordinates to values.

    data Matrix i b = Matrix
      { size :: Size i        -- ^ Dimensions of the matrix.
      , unM  :: [(MIx i, b)]  -- ^ Association of indices to values.
      }
      deriving (Eq, Ord, Functor, Foldable, Traversable)

    ------------------------------------------------------------------------
    -- * Operations and query on matrix size.
    ------------------------------------------------------------------------

    -- | 'True' iff the matrix is square.

    square :: Ix i => Matrix i b -> Bool
    square m = rows (size m) == cols (size m)

    -- | Returns 'True' iff the matrix is empty.

    isEmpty :: (Num i, Ix i) => Matrix i b -> Bool
    isEmpty m = rows sz <= 0 || cols sz <= 0
      where sz = size m

    -- | Compute the matrix size of the union of two matrices.

    supSize :: Ord i => Matrix i a -> Matrix i b -> Size i
    supSize (Matrix (Size r1 c1) _) (Matrix (Size r2 c2) _) =
      Size (max r1 r2) (max c1 c2)

    -- | Compute the matrix size of the intersection of two matrices.

    infSize :: Ord i => Matrix i a -> Matrix i b -> Size i
    infSize (Matrix (Size r1 c1) _) (Matrix (Size r2 c2) _) =
      Size (min r1 r2) (min c1 c2)

    ------------------------------------------------------------------------
    -- * Creating matrices and converting to lists.
    ------------------------------------------------------------------------

    -- | Constructs a matrix from a list of @(index, value)@-pairs.
    --   @O(n)@ where @n@ is size of the list.
    --
    --   Precondition: indices are unique.

    fromIndexList :: (Ord i, HasZero b) => Size i -> [(MIx i, b)] -> Matrix i b
    fromIndexList sz = Matrix sz
                     . List.sortBy (compare `on` fst)
                     . filter ((zeroElement /=) . snd)

    -- | @'fromLists' sz rs@ constructs a matrix from a list of lists of
    --   values (a list of rows).
    --   @O(size)@ where @size = rows × cols@.
    --
    --   Precondition:
    --   @'length' rs '==' 'rows' sz@ and
    --   @'all' (('cols' sz '==') . 'length') rs@.

    fromLists :: (Ord i, Num i, Enum i, HasZero b) => Size i -> [[b]] -> Matrix i b
    fromLists sz bs = fromIndexList sz $ zip ([ MIx i j | i <- [1..rows sz]
                                                        , j <- [1..cols sz]]) (concat bs)

    -- | Converts a sparse matrix to a sparse list of rows.
    --   @O(n)@ where @n@ is the number of non-zero entries of the matrix.
    --
    --   Only non-empty rows are generated.
    --

    toSparseRows :: (Eq i) => Matrix i b -> [(i,[(i,b)])]
    toSparseRows (Matrix _ []) = []
    toSparseRows (Matrix _ ((MIx i j, b) : m)) = aux i [(j,b)] m
      where aux i' []  [] = []
            aux i' row [] = [(i', reverse row)]
            aux i' row ((MIx i j, b) : m)
              | i' == i   = aux i' ((j,b):row) m
              | otherwise = (i', reverse row) : aux i [(j,b)] m

    -- | Turn a sparse vector into a vector by inserting a fixed element
    --   at the missing positions.
    --   @O(size)@ where @size@ is the dimension of the vector.

    blowUpSparseVec :: (Integral i) => b -> i -> [(i,b)] -> [b]
    blowUpSparseVec zero n l = aux 1 l
      where aux i []           = List.genericReplicate (n + 1 - i) zero
            aux i l@((j,b):l')
              | i > j || i > n = __IMPOSSIBLE__
              | i == j         = b    : aux (i + 1) l'
              | otherwise      = zero : aux (i + 1) l

    -- Older implementation without replicate.
    blowUpSparseVec' :: (Ord i, Num i, Enum i) => b -> i -> [(i,b)] -> [b]
    blowUpSparseVec' zero n l = aux 1 l
      where aux i [] | i > n = []
                     | otherwise = zero : aux (i+1) []
            aux i ((j,b):l) | i <= n && j == i = b : aux (succ i) l
            aux i ((j,b):l) | i <= n && j >= i = zero : aux (succ i) ((j,b):l)
            aux i l = __IMPOSSIBLE__
              -- error $ "blowUpSparseVec (n = " ++ show n ++ ") aux i=" ++ show i ++ " j=" ++ show (fst (head l)) ++ " length l = " ++ show (length l)

    -- | Converts a matrix to a list of row lists.
    --   @O(size)@ where @size = rows × cols@.

    toLists :: (Integral i, HasZero b) => Matrix i b -> [[b]]
    toLists m@(Matrix size@(Size nrows ncols) _) =
        blowUpSparseVec emptyRow nrows $
          map (mapSnd (blowUpSparseVec zeroElement ncols)) $ toSparseRows m
      where
        emptyRow = List.genericReplicate ncols zeroElement

    ------------------------------------------------------------------------
    -- * Combining and querying matrices
    ------------------------------------------------------------------------

    -- | Returns 'Just b' iff it is a 1x1 matrix with just one entry 'b'.
    --   @O(1)@.

    isSingleton :: (Eq i, Num i, HasZero b) => Matrix i b -> Maybe b
    isSingleton (Matrix (Size 1 1) [(_,b)]) = Just b
    isSingleton (Matrix (Size 1 1) []     ) = Just zeroElement
    isSingleton (Matrix (Size 1 1) _      ) = __IMPOSSIBLE__
    isSingleton _ = Nothing

    -- | @'diagonal' m@ extracts the diagonal of @m@.
    --
    --   For non-square matrices, the length of the diagonal is
    --   the minimum of the dimensions of the matrix.

    class Diagonal m e | m -> e where
      diagonal :: m -> [e]

    -- | Diagonal of sparse matrix.
    --
    --   @O(n)@ where @n@ is the number of non-zero elements in the matrix.

    instance (Integral i, HasZero b) => Diagonal (Matrix i b) b where
      diagonal (Matrix (Size r c) m) =
        blowUpSparseVec zeroElement (min r c) $
          mapMaybe (\ ((MIx i j), b) -> if i==j then Just (i,b) else Nothing) m

    -- | Transposable things.

    class Transpose a where
      transpose :: a -> a

    -- | Size of transposed matrix.

    instance Transpose (Size i) where
      transpose (Size n m) = Size m n

    -- | Transposing coordinates.

    instance Transpose (MIx i) where
      transpose (MIx i j) = MIx j i

    -- | Matrix transposition.
    --
    --   @O(n log n)@ where @n@ is the number of non-zero elements in the matrix.

    instance Ord i => Transpose (Matrix i b) where
      transpose (Matrix size m) =
        Matrix (transpose size) $
          List.sortBy (compare `on` fst) $
            map (mapFst transpose) m

    -- | General pointwise combination function for association lists.
    --   @O(n1 + n2)@ where @ni@ is the number of non-zero element in matrix @i@.
    --
    --   In @zipAssocWith fs gs f g h l l'@,
    --
    --   @fs@ is possibly more efficient version of
    --   @'mapMaybe' (\ (i, a) -> (i,) <$> f a)@, and same for @gs@ and @g@.

    zipAssocWith :: (Ord i)
      => ([(i,a)] -> [(i,c)]) -- ^ Only left map remaining.
      -> ([(i,b)] -> [(i,c)]) -- ^ Only right map remaining.
      -> (a -> Maybe c)       -- ^ Element only present in left map.
      -> (b -> Maybe c)       -- ^ Element only present in right map.
      -> (a -> b -> Maybe c)  -- ^ Element present in both maps.
      -> [(i,a)] -> [(i,b)] -> [(i,c)]
    zipAssocWith fs gs f g h = merge
      where
        merge m1 [] = mapMaybe (\ (i, a) -> (i,) <$> f a) m1
        merge [] m2 = mapMaybe (\ (i, b) -> (i,) <$> g b) m2
        merge m1@((i,a):m1') m2@((j,b):m2') =
          case compare i j of
            LT -> mcons ((i,) <$> f a)   $ merge m1' m2
            GT -> mcons ((j,) <$> g b)   $ merge m1  m2'
            EQ -> mcons ((i,) <$> h a b) $ merge m1' m2'

    -- | Instance of 'zipAssocWith' which keeps longer assoc lists.
    --   @O(n1 + n2)@.

    unionAssocWith  :: (Ord i)
      => (a -> Maybe c)       -- ^ Element only present in left map.
      -> (b -> Maybe c)       -- ^ Element only present in right map.
      -> (a -> b -> Maybe c)  -- ^ Element present in both maps.
      -> [(i,a)] -> [(i,b)] -> [(i,c)]
    unionAssocWith f g h = zipAssocWith (map_ f) (map_ g) f g h
      where
        map_ f = mapMaybe (\ (i, a) -> (i,) <$> f a)

    -- | General pointwise combination function for sparse matrices.
    --   @O(n1 + n2)@.

    zipMatrices :: forall a b c i . (Ord i)
      => (a -> c)       -- ^ Element only present in left matrix.
      -> (b -> c)       -- ^ Element only present in right matrix.
      -> (a -> b -> c)  -- ^ Element present in both matrices.
      -> (c -> Bool)    -- ^ Result counts as zero?
      -> Matrix i a -> Matrix i b -> Matrix i c
    zipMatrices f g h zero m1 m2 = Matrix (supSize m1 m2) $
      unionAssocWith (drop0 . f) (drop0 . g) (\ a -> drop0 . h a) (unM m1) (unM m2)
      where
        drop0 = filterMaybe (not . zero)

    -- | @'add' (+) m1 m2@ adds @m1@ and @m2@, using @(+)@ to add values.
    --   @O(n1 + n2)@.
    --
    --   Returns a matrix of size @'supSize' m1 m2@.

    add :: (Ord i, HasZero a) => (a -> a -> a) -> Matrix i a -> Matrix i a -> Matrix i a
    add plus = zipMatrices id id plus (== zeroElement)

    -- | @'intersectWith' f m1 m2@ build the pointwise conjunction @m1@ and @m2@.
    --   Uses @f@ to combine non-zero values.
    --   @O(n1 + n2)@.
    --
    --   Returns a matrix of size @infSize m1 m2@.

    intersectWith :: (Ord i) => (a -> a -> a) -> Matrix i a -> Matrix i a -> Matrix i a
    intersectWith f m1 m2 = Matrix (infSize m1 m2) $ interAssocWith f (unM m1) (unM m2)

    -- | Association list intersection.
    --   @O(n1 + n2)@.
    --
    --   @interAssocWith f l l' = { (i, f a b) | (i,a) ∈ l and (i,b) ∈ l' }@
    --
    --   Used to combine sparse matrices, it might introduce zero elements
    --   if @f@ can return zero for non-zero arguments.

    interAssocWith :: (Ord i) => (a -> a -> a) -> [(i,a)] -> [(i,a)] -> [(i,a)]
    interAssocWith f [] m = []
    interAssocWith f l [] = []
    interAssocWith f l@((i,a):l') m@((j,b):m')
        | i < j = interAssocWith f l' m
        | i > j = interAssocWith f l m'
        | otherwise = (i, f a b) : interAssocWith f l' m'

    -- | @'mul' semiring m1 m2@ multiplies matrices @m1@ and @m2@.
    --   Uses the operations of the semiring @semiring@ to perform the
    --   multiplication.
    --
    --   @O(n1 + n2 log n2 + Σ(i <= r1) Σ(j <= c2) d(i,j))@ where
    --   @r1@ is the number of non-empty rows in @m1@ and
    --   @c2@ is the number of non-empty columns in @m2@ and
    --   @d(i,j)@  is the bigger one of the following two quantifies:
    --   the length of sparse row @i@ in @m1@ and
    --   the length of sparse column @j@ in @m2@.
    --
    --   Given dimensions @m1 : r1 × c1@ and @m2 : r2 × c2@,
    --   a matrix of size @r1 × c2@ is returned.
    --   It is not necessary that @c1 == r2@, the matrices are implicitly
    --   patched with zeros to match up for multiplication.
    --   For sparse matrices, this patching is a no-op.

    mul :: (Ix i, Eq a)
        => Semiring a -> Matrix i a -> Matrix i a -> Matrix i a
    mul semiring m1 m2 = Matrix (Size { rows = rows (size m1), cols = cols (size m2) }) $
      [ (MIx i j, b)
        | (i,v) <- toSparseRows m1
        , (j,w) <- toSparseRows $ transpose m2
        , let b = inner v w
        , b /= zero
      ]
      where
         zero  = Semiring.zero semiring
         plus  = Semiring.add  semiring
         times = Semiring.mul  semiring
         inner v w = List.foldl' plus zero $
                       map snd $ interAssocWith times v w

    -- | Pointwise comparison.
    --   Only matrices with the same dimension are comparable.
    instance (Ord i, PartialOrd a) => PartialOrd (Matrix i a) where
      comparable m n
        | size m /= size n = POAny
        | otherwise        = Fold.fold $
                               zipMatrices onlym onlyn both trivial m n
        where
          -- If an element is only in @m@, then its 'Unknown' in @n@
          -- so it gotten better at best, in any case, not worse.
          onlym o = POGT
          -- If an element is only in @n@, then its 'Unknown' in @m@
          -- so we have strictly less information.
          onlyn o = POLT
          both    = comparable
          -- The zero element of the result sparse matrix is the
          -- neutral element of the monoid.
          trivial = (==mempty)

    ------------------------------------------------------------------------
    -- Modifying matrices

    -- | @'addColumn' x m@ adds a new column to @m@, after the columns
    -- already existing in the matrix. All elements in the new column get
    -- set to @x@.

    addColumn :: (Num i, HasZero b) => b -> Matrix i b -> Matrix i b
    addColumn x m | x == zeroElement = m { size = (size m) { cols = cols (size m) + 1 }}
                  | otherwise = __IMPOSSIBLE__

    -- | @'addRow' x m@ adds a new row to @m@, after the rows already
    -- existing in the matrix. All elements in the new row get set to @x@.

    addRow :: (Num i, HasZero b) => b -> Matrix i b -> Matrix i b
    addRow x m | x == zeroElement = m { size = (size m) { rows = rows (size m) + 1 }}
               | otherwise = __IMPOSSIBLE__

    ------------------------------------------------------------------------
    -- * Printing
    ------------------------------------------------------------------------

    instance (Integral i, HasZero b, Show i, Show b) => Show (Matrix i b) where
      showsPrec _ m =
        showString "Agda.Termination.SparseMatrix.fromLists " . shows (size m) .
        showString " " . shows (toLists m)

    instance (Integral i, HasZero b, Pretty b) =>
             Pretty (Matrix i b) where
    --  pretty = vcat . map (hsep . map pretty) . toLists
      pretty = vcat
             . map text
             . lines
             . Boxes.render
             . Boxes.hsep 1 Boxes.right
             . map ( Boxes.vcat Boxes.right
                   . map ( Boxes.alignHoriz Boxes.right 4
                         . Boxes.text . render . pretty
                         )
                   )
             . toLists
             . transpose
    -- ADAPTED FROM:
    -- http://www.tedreed.info/programming/2012/06/02/how-to-use-textprettyprintboxes/
    -- print_table :: [[String]] -> IO ()
    -- print_table rows = printBox $ hsep 2 left (map (vcat left . map text) (transpose rows))