{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DeriveDataTypeable    #-}

module Agda.Utils.Permutation where

import Prelude hiding (drop, null)

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

import Data.Foldable (Foldable)
import Data.Traversable (Traversable)

import Data.Data (Data)

import Agda.Syntax.Position (KillRange(..))
import Agda.Utils.Functor
import Agda.Utils.List ((!!!))
import Agda.Utils.Null
import Agda.Utils.Size

#include "undefined.h"
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)

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 KillRange Permutation where
killRange = id

-- | @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.findIndex (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 (flip const) 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 = fmap (Perm (size xs)) \$ topo g
where
nodes     = zip [0..] xs
g         = [ (n, parents x) | (n, x) <- nodes ]
parents x = [ n | (n, y) <- nodes, parent y x ]

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)

instance KillRange a => KillRange (Drop a) where
killRange = fmap killRange

-- | 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