-- | Lexicographic order search, more or less as defined in
--      \"A Predicative Analysis of Structural Recursion\" by
--      Andreas Abel and Thorsten Altenkirch.

module Agda.Termination.Lexicographic
  ( LexOrder
  , RecBehaviour(..)
  , Column
  , recBehaviourInvariant
  , fromDiagonals
  , lexOrder
  , Agda.Termination.Lexicographic.tests
  ) where

import Agda.Utils.QuickCheck
import Agda.Termination.Matrix (Size (..))
import qualified Agda.Termination.Matrix as M
import Agda.Utils.Either
import Agda.Utils.List
import Agda.Utils.TestHelpers
import Agda.Termination.CallGraph
import Control.Arrow
import Data.List
import Data.Map (Map, (!))
import qualified Data.Map as Map
import Data.Array (Array, Ix)
import qualified Data.Array as Array
import Data.Set (Set)
import qualified Data.Set as Set

-- | A lexicographic ordering for the recursion behaviour of a
-- given function is a permutation of the argument indices which can
-- be used to show that the function terminates. See the paper
-- referred to above for more details.

type LexOrder arg = [arg]

-- | A recursion behaviour expresses how a certain function calls
-- itself (transitively). For every argument position there is a value
-- ('Column') describing how the function calls itself for that
-- particular argument. See also 'recBehaviourInvariant'.

data RecBehaviour arg call =
  RB { columns :: Map arg (Column call)
     , calls   :: Set call
       -- ^ The indices to the columns.
     , size    :: Size Integer
     }
  deriving Show

-- | A column expresses how the size of a certain argument changes in
-- the various recursive calls a function makes to itself
-- (transitively).

type Column call = Map call Order

-- | 'RecBehaviour' invariant: the size must match the real size of
-- the recursion behaviour, and all columns must have the same
-- indices.

recBehaviourInvariant :: Eq call => RecBehaviour arg call -> Bool
recBehaviourInvariant rb =
  genericLength (Map.elems $ columns rb) == cols (size rb)
  &&
  all (== rows (size rb))
      (map (genericLength . Map.elems) $ Map.elems $ columns rb)
  &&
  allEqual (calls rb : (map Map.keysSet $ Map.elems $ columns rb))

-- Generates a recursion behaviour.

instance (Arbitrary call, Arbitrary arg, Ord arg, Ord call)
  => Arbitrary (RecBehaviour call arg) where
  arbitrary = do
    calls <- fmap nub $ listOf arbitrary
    args  <- fmap nub $ listOf arbitrary
    let rows = genericLength calls
        cols = genericLength args
        sz   = Size { rows = rows, cols = cols }
        colGen = do
          col <- vectorOf (fromIntegral rows) arbitrary
          return $ Map.fromList (zip calls col)
    cols <- fmap (zip args) $ vectorOf (fromIntegral cols) colGen
    return $ RB { columns = Map.fromList cols
                , calls   = Set.fromList calls
                , size    = sz }

instance (CoArbitrary call, CoArbitrary arg) => CoArbitrary (RecBehaviour call arg) where
  coarbitrary (RB c cs s) =
    coarbitrary (map (id *** Map.toList) $ Map.toList c) .
    coarbitrary (Set.toList cs) .
    coarbitrary s

prop_recBehaviour_Arbitrary :: RecBehaviour Integer Integer -> Bool
prop_recBehaviour_Arbitrary = recBehaviourInvariant

-- | Checks whether there are any calls left in the recursion
-- behaviour.

noCallsLeft :: RecBehaviour arg call -> Bool
noCallsLeft rb = rows (size rb) == 0

-- | Constructs a recursion behaviour from a list of matrix diagonals
-- (\"rows\"). Note that the @call@ indices do not need to be
-- distinct, since they are paired up with unique 'Integer's.
--
-- Precondition: all arrays should have the same bounds.

fromDiagonals :: (Ord call, Ix arg)
              => [(call, Array arg Order)] -> RecBehaviour arg (Integer, call)
fromDiagonals []   = RB { columns = Map.fromList []
                        , calls   = Set.empty
                        , size    = Size 0 0 }
fromDiagonals rows = RB { columns = Map.fromList $ zip args cols
                        , calls   = Set.fromList calls
                        , size    = Size { rows = genericLength rows
                                         , cols = genericLength cols }
                        }
  where
  calls = zip [1 ..] $ map fst rows
  cols = map Map.fromList $ map (zip calls) $ transpose $
         map (Array.elems . snd) rows
  args = Array.range $ Array.bounds $ snd $ head rows

prop_fromDiagonals m =
  forAll (vectorOf (fromIntegral $ rows $ M.size m) arbitrary) $ \calls ->
    let oss = zip calls $
              map (Array.listArray (1, cols $ M.size m)) $
              M.toLists m
        rb = fromDiagonals oss :: RecBehaviour Integer (Integer, Bool)
    in
    recBehaviourInvariant rb
    &&
    if rows (M.size m) == 0 then
      rows (size rb) == 0
     else
      size rb == M.size m

-- | Checks if this \"column\" is well-behaved (all calls decreasing,
-- at least one strictly decreasing).

okColumn :: Column call -> Bool
okColumn col = any (== Lt) col' && all (/= Unknown) col'
  where col' = Map.elems col

-- | @'newBehaviour' n rb@ computes a new recursion behaviour from
-- @rb@ by removing all \"rows\" (calls) for which the @n@-th element
-- is 'Lt', and also completely removing the @n@-th column.
--
-- Precondition: there has to be an @n@-th column.

newBehaviour :: (Ord arg, Ord call)
             => arg -> RecBehaviour arg call -> RecBehaviour arg call
newBehaviour n rb =
  RB { columns = Map.map remove $ Map.delete n $ columns rb
     , calls   = Set.difference (calls rb)
                                (Set.fromList indicesToRemove)
     , size    = Size { rows = rows (size rb) -
                               genericLength indicesToRemove
                      , cols = cols (size rb) - 1 }
     }
  where
  Just colN       = Map.lookup n $ columns rb
  indicesToRemove = map fst $ filter ((== Lt) . snd) $ Map.toList colN
  remove colJ     = foldr Map.delete colJ indicesToRemove

prop_newBehaviour :: RecBehaviour Integer Integer -> Property
prop_newBehaviour rb =
  not (cols (size rb) == 0) ==>
    forAll (elements $ Map.keys $ columns rb) $ \n ->
      recBehaviourInvariant (newBehaviour n rb)

-- | @'correctLexOrder' rs ps@ checks that the permutation @ps@ really
-- induces a lexicographic ordering which shows that the function
-- represented by the recursion behaviour @rs@ terminates.

correctLexOrder :: (Ord arg, Ord call)
                => RecBehaviour arg call -> LexOrder arg -> Bool
correctLexOrder rb []        = noCallsLeft rb
correctLexOrder rb (p0 : ps) =
  okColumn (columns rb ! p0) && correctLexOrder (newBehaviour p0 rb) ps

-- | Tries to compute a lexicographic ordering for the given recursion
-- behaviour. This algorithm should be complete.
--
-- If no lexicographic ordering can be found, then two sets are
-- returned:
--
-- * A set of argument positions which are not properly decreasing, and
--
-- * the calls where these problems show up.

lexOrder :: (Ord arg, Ord call) =>
  RecBehaviour arg call -> Either (Set arg, Set call) (LexOrder arg)
lexOrder rb | noCallsLeft rb = Right []
            | otherwise      = case okColumns of
  []      -> Left (Map.keysSet $ columns rb, calls rb)
  (n : _) -> case lexOrder (newBehaviour n rb) of
    Left err -> Left err
    Right ps -> Right $ n : ps
  where
  okColumns = map fst $ filter snd $
              map (id *** okColumn) $
              Map.toList $ columns rb

prop_lexOrder :: RecBehaviour Integer Integer -> Property
prop_lexOrder rb =
  let mPerm = lexOrder rb
      Right perm = mPerm
  in
  isRight mPerm ==>
    classify (cols (size rb) >= 2) "interesting" $
    correctLexOrder rb perm

prop_lexOrder_noArgs =
  forAll positive $ \n ->
    isLeft (lexOrder $ rb n)
    where rb :: Integer -> RecBehaviour Integer Integer
          rb n = RB { columns = Map.empty
                    , calls   = Set.fromList [1 .. n]
                    , size    = Size { rows = n, cols = 0 }
                    }

------------------------------------------------------------------------
-- All tests

tests :: IO Bool
tests = runTests "Agda.Termination.Lexicographic"
  [ quickCheck' prop_recBehaviour_Arbitrary
  , quickCheck' prop_fromDiagonals
  , quickCheck' prop_newBehaviour
  , quickCheckWith' Args{ replay     = Nothing
                        , maxSuccess = 50
                        , maxDiscard = 200
                        , maxSize    = 20
                        }
                    prop_lexOrder
  , quickCheck' prop_lexOrder_noArgs
  ]