{-# LANGUAGE TemplateHaskell #-}

module Internal.Termination.SparseMatrix
  ( matrix
  , tests
  ) where

import Agda.Termination.Semiring (HasZero(..), Semiring)
import qualified Agda.Termination.Semiring as Semiring
import Agda.Termination.SparseMatrix
import Agda.Utils.Functor
import Agda.Utils.Tuple

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

import Internal.Helpers

------------------------------------------------------------------------
-- * Generators, properties and tests
------------------------------------------------------------------------

-- ** Size
------------------------------------------------------------------------

instance Integral i => Arbitrary (Size i) where
  arbitrary = do
    r <- natural
    c <- natural
    return $ Size { rows = fromInteger r, cols = fromInteger c }

instance CoArbitrary i => CoArbitrary (Size i) where
  coarbitrary (Size rs cs) = coarbitrary rs . coarbitrary cs

-- | Size invariant: dimensions are non-negative.

sizeInvariant :: (Ord i, Num i) => Size i -> Bool
sizeInvariant sz = rows sz >= 0 && cols sz >= 0

prop_Arbitrary_Size :: Size Integer -> Bool
prop_Arbitrary_Size = sizeInvariant

prop_size :: TM -> Bool
prop_size m = sizeInvariant (size m)

-- ** Matrix indices
------------------------------------------------------------------------

instance Integral i => Arbitrary (MIx i) where
  arbitrary = MIx <$> positive <*> positive

instance CoArbitrary i => CoArbitrary (MIx i) where
  coarbitrary (MIx r c) = coarbitrary r . coarbitrary c

-- | Indices must be positive, @>= 1@.

mIxInvariant :: (Ord i, Num i) => MIx i -> Bool
mIxInvariant i = row i >= 1 && col i >= 1

prop_Arbitrary_MIx :: MIx Integer -> Bool
prop_Arbitrary_MIx = mIxInvariant

-- ** Matrices
------------------------------------------------------------------------

-- | Matrix indices are lexicographically sorted with no duplicates.
--   All indices must be within bounds.
matrixInvariant :: (Num i, Ix i, HasZero b) => Matrix i b -> Bool
matrixInvariant (Matrix size@Size{ rows, cols} m) =
  sizeInvariant size
  && all inBounds m
  && all nonZero  m
  && strictlySorted (MIx 0 0) m
  where
    inBounds (MIx i j, _) = 1 <= i && i <= rows
                         && 1 <= j && j <= cols
    nonZero (_, b) = b /= zeroElement

-- | Check whether an association list is ordered and
--   deterministic, a partial function from @i@ to @b@.
strictlySorted :: (Ord i) => i -> [(i, b)] -> Bool
strictlySorted i []            = True
strictlySorted i ((i', _) : l) = i < i' && strictlySorted i' l
-- Ord MIx should be the lexicographic order already (Haskell report).

-- | Generates a matrix of the given size, using the given generator
-- to generate the rows.

matrixUsingRowGen :: (Integral i, HasZero b)
  => Size i
  -> (i -> Gen [b])
     -- ^ The generator is parameterised on the size of the row.
  -> Gen (Matrix i b)
matrixUsingRowGen sz rowGen = do
  rows <- vectorOf (fromIntegral $ rows sz) (rowGen $ cols sz)
  return $ fromLists sz rows

-- | Generates a matrix of the given size.

matrix :: (Integral i, Arbitrary b, HasZero b)
       => Size i -> Gen (Matrix i b)
matrix sz = matrixUsingRowGen sz (\n -> vectorOf (fromIntegral n) arbitrary)

prop_matrix :: Size Int -> Property
prop_matrix sz = forAll (matrix sz :: Gen TM) $ \ m -> size m == sz

-- | Generate a matrix of arbitrary size.

instance (Integral i, Arbitrary b, HasZero b)
         => Arbitrary (Matrix i b) where
  arbitrary     = matrix =<< arbitrary

instance (Integral i, CoArbitrary b, HasZero b) => CoArbitrary (Matrix i b) where
  coarbitrary m = coarbitrary (toLists m)

-- | This matrix type is used for tests.

type TM = Matrix Int Int

prop_Arbitrary_Matrix :: TM -> Bool
prop_Arbitrary_Matrix = matrixInvariant

-- ** Matrix operations

-- | 'fromIndexList' is identity on well-formed sparse matrices.

prop_fromIndexList :: TM -> Bool
prop_fromIndexList m@(Matrix size vs) = fromIndexList size vs == m

-- | Converting a matrix to a list of lists and back is the identity.

prop_fromLists_toLists :: TM -> Bool
prop_fromLists_toLists m = fromLists (size m) (toLists m) == m

-- | Any 1x1 matrix is a singleton.

prop_isSingleton :: Int -> Bool
prop_isSingleton b = Just b == (isSingleton (fromLists (Size 1 1) [[b]] :: TM))

-- | The length of the diagonal is the minimum of the number of
--   rows and columns of the matrix.

prop_diagonal :: TM -> Bool
prop_diagonal m@(Matrix (Size r c) _) =
    length (diagonal m) == min r c

-- prop_diagonal' n =
--   forAll natural $ \n ->
--   forAll (matrix (Size n n) :: Gen TM) $ \m ->
--     length (diagonal m) == n

-- | Transposing twice is the identity.

prop_transpose :: TM -> Bool
prop_transpose m = matrixInvariant m' && m == transpose m'
  where m' = transpose m

-- | Verify 'toSparseRows' against an alternative implementation which
-- serves as specification.

toSparseRows' :: (Eq i) => Matrix i b -> [(i,[(i,b)])]
toSparseRows' (Matrix _ m) =
  -- group list by row index
  for (List.groupBy ((==) `on` (row . fst)) m) $ \ ((MIx i j, b) : vs) ->
  -- turn each group into a sparse row
    (i, (j,b) : map (mapFst col) vs)

prop_toSparseRows :: TM -> Bool
prop_toSparseRows m = toSparseRows m == toSparseRows' m

prop_addColumn :: TM -> Bool
prop_addColumn m =
  matrixInvariant m'
  &&
  map init (toLists m') == toLists m
  where
  m' = addColumn zeroElement m


prop_addRow :: TM -> Bool
prop_addRow m =
  matrixInvariant m'
  &&
  init (toLists m') == toLists m
  where
  m' = addRow zeroElement m


-- ** Matrix addition


-- | Old implementation of 'zipMatrices'.

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) (merge (unM m1) (unM m2))
  where
    merge :: [(MIx i,a)] -> [(MIx i,b)] -> [(MIx i,c)]
    merge [] m2 = filter (not . zero . snd) $ map (mapSnd g) m2
    merge m1 [] = filter (not . zero . snd) $ map (mapSnd f) m1
    merge m1@((i,a):m1') m2@((j,b):m2') =
      case compare i j of
        LT -> if zero c then r else (i,c) : r where c = f a   ; r = merge m1' m2
        GT -> if zero c then r else (j,c) : r where c = g b   ; r = merge m1  m2'
        EQ -> if zero c then r else (i,c) : r where c = h a b ; r = merge m1' m2'

-- | Verify 'zipMatrices' against older implementation

prop_zipMatrices_correct :: TM -> TM -> Bool
prop_zipMatrices_correct m1 m2 =
  zipMatrices id id (+) (== 0) m1 m2 == zipMatrices' id id (+) (== 0) m1 m2


-- | Matrix addition is well-defined, associative and commutative.

prop_add :: Size Int -> Property
prop_add sz =
  forAll (three (matrix sz :: Gen TM)) $ \(m1, m2, m3) ->
    let m' = add (+) m1 m2 in
      isAssociative (add (+)) m1 m2 m3 &&
      isCommutative (add (+)) m1 m2 &&
      matrixInvariant m' &&
      size m' == size m1

-- | Verify addition against an older implementation.
--
--   The older implementation did not fully preserve sparsity,
--   i.e., introduced zeros.  Thus, we need to convert to lists to
--   obtain equal results.

prop_add_correct :: TM -> TM -> Bool
prop_add_correct m1 m2 = toLists (add (+) m1 m2) == toLists (add' (+) m1 m2)
  where
    add' :: (Ord i) => (a -> a -> a) -> Matrix i a -> Matrix i a -> Matrix i a
    add' plus m1 m2 = Matrix (supSize m1 m2) $ mergeAssocWith plus (unM m1) (unM m2)
      where
        mergeAssocWith :: (Ord i) => (a -> a -> a) -> [(i,a)] -> [(i,a)] -> [(i,a)]
        mergeAssocWith f [] m = m
        mergeAssocWith f l [] = l
        mergeAssocWith f l@((i,a):l') m@((j,b):m')
            | i < j = (i,a) : mergeAssocWith f l' m
            | i > j = (j,b) : mergeAssocWith f l m'
            | otherwise = (i, f a b) : mergeAssocWith f l' m'

-- ** Matrix multiplication

-- | Specification of 'interAssocWith'.

interAssocWith' :: (Eq i) => (a -> a -> a) -> [(i,a)] -> [(i,a)] -> [(i,a)]
interAssocWith' f l l' = [ (i, f a b) | (i,a) <- l, (j,b) <- l', i == j ]

-- | Efficient implementation of 'interAssocWith' matches its specification.

no_tested_prop_interAssocWith_correct :: [(Int,Int)] -> [(Int,Int)] -> Bool
no_tested_prop_interAssocWith_correct xs ys =
  interAssocWith (*) l l' == interAssocWith' (*) l l'
  where
    l  = List.sortBy (compare `on` fst) xs
    l' = List.sortBy (compare `on` fst) ys

interAssocWith2 :: Ord i => (a -> a -> a) -> [(i,a)] -> [(i,a)] -> [(i,a)]
interAssocWith2 f = zipAssocWith (const []) (const []) (const Nothing) (const Nothing) (\ a -> Just . f a)

no_tested_prop_interAssocWith_correct2 :: [(Int,Int)] -> [(Int,Int)] -> Bool
no_tested_prop_interAssocWith_correct2 xs ys =
  interAssocWith (*) xs ys == interAssocWith2 (*) xs ys
  where
    l  = List.sortBy (compare `on` fst) xs
    l' = List.sortBy (compare `on` fst) ys

-- | Matrix multiplication is well-defined and associative.

prop_mul :: Size Int -> Property
prop_mul sz =
  mapSize (`div` 2) $
  forAll (two natural) $ \(c2, c3) ->
  forAll (matrix sz :: Gen TM) $ \m1 ->
  forAll (matrix (Size { rows = cols sz, cols = c2 })) $ \m2 ->
  forAll (matrix (Size { rows = c2, cols = c3 })) $ \m3 ->
    let m' = mult m1 m2 in
      isAssociative mult m1 m2 m3 &&
      matrixInvariant m' &&
      size m' == Size { rows = rows sz, cols = c2 }
  where mult = mul Semiring.intSemiring

------------------------------------------------------------------------
-- * All tests
------------------------------------------------------------------------

-- Template Haskell hack to make the following $allProperties work
-- under ghc-7.8.
return [] -- KEEP!

-- | All tests as collected by 'allProperties'.
--
-- Using 'allProperties' is convenient and superior to the manual
-- enumeration of tests, since the name of the property is added
-- automatically.

tests :: TestTree
tests = testProperties "Internal.Termination.SparseMatrix" $allProperties

