
module Agda.Utils.Permutation where

import Prelude hiding (drop, null)

import Control.DeepSeq
import Control.Monad (filterM)

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Functor.Identity
import qualified Data.List as List
import Data.Maybe
import Data.Array

import Data.Data (Data)

import GHC.Generics (Generic)

import Agda.Utils.Functor
import Agda.Utils.List ((!!!))
import Agda.Utils.Null
import Agda.Utils.Size
import Agda.Utils.Tuple

import Agda.Utils.Impossible

-- | Partial permutations. Examples:
--
--   @permute [1,2,0]   [x0,x1,x2] = [x1,x2,x0]@     (proper permutation).
--
--   @permute [1,0]     [x0,x1,x2] = [x1,x0]@        (partial permuation).
--
--   @permute [1,0,1,2] [x0,x1,x2] = [x1,x0,x1,x2]@  (not a permutation because not invertible).
--
--   Agda typing would be:
--   @Perm : {m : Nat}(n : Nat) -> Vec (Fin n) m -> Permutation@
--   @m@ is the 'size' of the permutation.
data Permutation = Perm { permRange :: Int, permPicks :: [Int] }
  deriving (Eq, Data, Generic)

instance Show Permutation where
  show (Perm n xs) = showx [0..n - 1] ++ " -> " ++ showx xs
    where showx = showList "," (\ i -> "x" ++ show i)
          showList :: String -> (a -> String) -> [a] -> String
          showList sep f [] = ""
          showList sep f [e] = f e
          showList sep f (e:es) = f e ++ sep ++ showList sep f es

instance Sized Permutation where
  size (Perm _ xs) = size xs

instance Null Permutation where
  empty = Perm 0 []
  null (Perm _ picks) = null picks

instance NFData Permutation

-- | @permute [1,2,0] [x0,x1,x2] = [x1,x2,x0]@
--   More precisely, @permute indices list = sublist@, generates @sublist@
--   from @list@ by picking the elements of list as indicated by @indices@.
--   @permute [1,3,0] [x0,x1,x2,x3] = [x1,x3,x0]@
--
--   Agda typing:
--   @permute (Perm {m} n is) : Vec A m -> Vec A n@
permute :: Permutation -> [a] -> [a]
permute p xs = map (fromMaybe __IMPOSSIBLE__) (safePermute p xs)

safePermute :: Permutation -> [a] -> [Maybe a]
safePermute (Perm _ is) xs = map (xs !!!!) is
  where
    xs !!!! n | n < 0     = Nothing
              | otherwise = xs !!! n

-- |  Invert a Permutation on a partial finite int map.
-- @inversePermute perm f = f'@
-- such that @permute perm f' = f@
--
-- Example, with map represented as @[Maybe a]@:
-- @
--   f    = [Nothing, Just a, Just b ]
--   perm = Perm 4 [3,0,2]
--   f'   = [ Just a , Nothing , Just b , Nothing ]
-- @
-- Zipping @perm@ with @f@ gives @[(0,a),(2,b)]@, after compression
-- with @catMaybes@.  This is an @IntMap@ which can easily
-- written out into a substitution again.

class InversePermute a b where
  inversePermute :: Permutation -> a -> b

instance InversePermute [Maybe a] [(Int,a)] where
  inversePermute (Perm n is) = catMaybes . zipWith (\ i ma -> (i,) <$> ma) is

instance InversePermute [Maybe a] (IntMap a) where
  inversePermute p = IntMap.fromList . inversePermute p

instance InversePermute [Maybe a] [Maybe a] where
  inversePermute p@(Perm n _) = tabulate . inversePermute p
    where tabulate m = for [0..n-1] $ \ i -> IntMap.lookup i m

instance InversePermute (Int -> a) [Maybe a] where
  inversePermute (Perm n xs) f = for [0..n-1] $ \ x -> f <$> List.elemIndex x xs

-- | Identity permutation.
idP :: Int -> Permutation
idP n = Perm n [0..n - 1]

-- | Restrict a permutation to work on @n@ elements, discarding picks @>=n@.
takeP :: Int -> Permutation -> Permutation
takeP n (Perm m xs) = Perm n $ filter (< n) xs

-- | Pick the elements that are not picked by the permutation.
droppedP :: Permutation -> Permutation
droppedP (Perm n xs) = Perm n $ [0..n-1] List.\\ xs

-- | @liftP k@ takes a @Perm {m} n@ to a @Perm {m+k} (n+k)@.
--   Analogous to 'Agda.TypeChecking.Substitution.liftS',
--   but Permutations operate on de Bruijn LEVELS, not indices.
liftP :: Int -> Permutation -> Permutation
liftP n (Perm m xs) = Perm (n + m) $ xs ++ [m..m+n-1]
-- liftP n (Perm m xs) = Perm (n + m) $ [0..n-1] ++ map (n+) xs -- WRONG, works for indices, but not for levels

-- | @permute (compose p1 p2) == permute p1 . permute p2@
composeP :: Permutation -> Permutation -> Permutation
composeP p1 (Perm n xs) = Perm n $ permute p1 xs
  {- proof:
      permute (compose (Perm xs) (Perm ys)) zs
      == permute (Perm (permute (Perm xs) ys)) zs
      == map (zs !!) (permute (Perm xs) ys)
      == map (zs !!) (map (ys !!) xs)
      == map (zs !! . ys !!) xs
      == map (\x -> zs !! (ys !! x)) xs
      == map (\x -> map (zs !!) ys !! x) xs  {- map f xs !! n == f (xs !! n) -}
      == map (map (zs !!) ys !!) xs
      == permute (Perm xs) (permute (Perm ys) zs)
  -}

-- | @invertP err p@ is the inverse of @p@ where defined,
--   otherwise defaults to @err@.
--   @composeP p (invertP err p) == p@
invertP :: Int -> Permutation -> Permutation
invertP err p@(Perm n xs) = Perm (size xs) $ elems tmpArray
  where tmpArray = accumArray (const id) err (0, n-1) $ zip xs [0..]

-- | Turn a possible non-surjective permutation into a surjective permutation.
compactP :: Permutation -> Permutation
compactP (Perm n xs) = Perm m $ map adjust xs
  where
    m            = List.genericLength xs
    missing      = [0..n - 1] List.\\ xs
    holesBelow k = List.genericLength $ filter (< k) missing
    adjust k = k - holesBelow k

-- | @permute (reverseP p) xs ==
--    reverse $ permute p $ reverse xs@
--
--   Example:
--   @
--        permute (reverseP (Perm 4 [1,3,0])) [x0,x1,x2,x3]
--     == permute (Perm 4 $ map (3-) [0,3,1]) [x0,x1,x2,x3]
--     == permute (Perm 4 [3,0,2]) [x0,x1,x2,x3]
--     == [x3,x0,x2]
--     == reverse [x2,x0,x3]
--     == reverse $ permute (Perm 4 [1,3,0]) [x3,x2,x1,x0]
--     == reverse $ permute (Perm 4 [1,3,0]) $ reverse [x0,x1,x2,x3]
--   @
--
--   With @reverseP@, you can convert a permutation on de Bruijn indices
--   to one on de Bruijn levels, and vice versa.
reverseP :: Permutation -> Permutation
reverseP (Perm n xs) = Perm n $ map ((n - 1) -) $ reverse xs
                  -- = flipP $ Perm n $ reverse xs

-- | @permPicks (flipP p) = permute p (downFrom (permRange p))@
--   or
--   @permute (flipP (Perm n xs)) [0..n-1] = permute (Perm n xs) (downFrom n)@
--
--   Can be use to turn a permutation from (de Bruijn) levels to levels
--   to one from levels to indices.
--
--   See 'Agda.Syntax.Internal.Patterns.numberPatVars'.
flipP :: Permutation -> Permutation
flipP (Perm n xs) = Perm n $ map (n - 1 -) xs

-- | @expandP i n π@ in the domain of @π@ replace the /i/th element by /n/ elements.
expandP :: Int -> Int -> Permutation -> Permutation
expandP i n (Perm m xs) = Perm (m + n - 1) $ concatMap expand xs
  where
    expand j
      | j == i    = [i..i + n - 1]
      | j < i     = [j]
      | otherwise = [j + n - 1]

-- | Stable topologic sort. The first argument decides whether its first
--   argument is an immediate parent to its second argument.
topoSort :: (a -> a -> Bool) -> [a] -> Maybe Permutation
topoSort parent xs = runIdentity $ topoSortM (\x y -> Identity $ parent x y) xs

topoSortM :: Monad m => (a -> a -> m Bool) -> [a] -> m (Maybe Permutation)
topoSortM parent xs = do
  let nodes     = zip [0..] xs
      parents x = map fst <$> filterM (\(_, y) -> parent y x) nodes
  g <- mapM (mapSndM parents) nodes
  return $ Perm (size xs) <$> topo g
  where

    topo :: Eq node => [(node, [node])] -> Maybe [node]
    topo [] = return []
    topo g  = case xs of
      []    -> fail "cycle detected"
      x : _ -> do
        ys <- topo $ remove x g
        return $ x : ys
      where
        xs = [ x | (x, []) <- g ]
        remove x g = [ (y, filter (/= x) ys) | (y, ys) <- g, x /= y ]

------------------------------------------------------------------------
-- * Drop (apply) and undrop (abstract)
------------------------------------------------------------------------

-- | Delayed dropping which allows undropping.
data Drop a = Drop
  { dropN    :: Int  -- ^ Non-negative number of things to drop.
  , dropFrom :: a    -- ^ Where to drop from.
  }
  deriving (Eq, Ord, Show, Data, Functor, Foldable, Traversable)

-- | Things that support delayed dropping.
class DoDrop a where

  doDrop :: Drop a -> a              -- ^ Perform the dropping.

  dropMore :: Int -> Drop a -> Drop a  -- ^ Drop more.
  dropMore n (Drop m xs) = Drop (m + n) xs

  unDrop :: Int -> Drop a -> Drop a  -- ^ Pick up dropped stuff.
  unDrop n (Drop m xs)
    | n <= m    = Drop (m - n) xs
    | otherwise = __IMPOSSIBLE__

instance DoDrop [a] where
  doDrop   (Drop m xs) = List.drop m xs

instance DoDrop Permutation where
  doDrop (Drop k (Perm n xs)) =
    Perm (n + m) $ [0..m-1] ++ map (+ m) (List.drop k xs)
    where m = -k
  unDrop m = dropMore (-m) -- allow picking up more than dropped
