{-  CAO Compiler
    Copyright (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. -}

{- |
Module      :  $Header$
Description :  Useful functions used in other modules.
Copyright   :  (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho
License     :  GPL

Maintainer  :  Paulo Silva <paufil@di.uminho.pt>
Stability   :  experimental
Portability :  portable

Useful functions used in other modules.
-}

module Language.CAO.Common.Utils where

import Control.Monad
import Data.DList (DList)
import qualified Data.DList as DL

import Data.List (
    genericIndex, genericTake, genericDrop, genericSplitAt, elemIndex, foldl')
import Data.Maybe (fromMaybe)

{-# INLINE singleton #-}
singleton :: a -> [a]
singleton = (:[])

{-# INLINE split #-}
split :: (a -> b) -> (a -> c) -> a -> (b, c)
split f g a = (f a, g a)

{-# INLINE mapPair #-}
mapPair :: (a -> b) -> (c -> d) -> (a,c) -> (b, d)
mapPair f g (a, c) = (f a, g c)

{-# INLINE mapFst #-}
mapFst :: (a -> b) -> (a, c) -> (b, c)
mapFst f = mapPair f id

{-# INLINE mapSnd #-}
mapSnd :: (c -> d) -> (a, c) -> (a, d)
mapSnd = mapPair id

{-# INLINE swap #-}
swap :: (a, b) -> (b, a)
swap (a, b) = (b, a)

swaps :: [(a, (b, c))] -> [(b, a, c)]
swaps = map (\ (a, (b, c)) -> (b, a, c))

initLast :: [a] -> ([a], a)
initLast [] = error "<initLast>: unexpected emtpy case"
initLast [x] = ([], x)
initLast (x:xs) = let
        (i, l) = initLast xs
    in (x:i, l)

(.$.) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
(f .$. g) a b = f (g a b)

apM :: Monad m => m (a -> b) -> a -> m b
apM f a = f >>= \ f' -> return $ f' a


{-# INLINE mapAndUnzip3M #-}
mapAndUnzip3M      :: (Monad m) => (a -> m (b, c, d)) -> [a] -> m ([b], [c], [d])
mapAndUnzip3M f xs =  fold3M f (:) (:) (:) ([], [], []) xs

{-# INLINE concatMapM #-}
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f = foldMM f (++) [] 

{-# INLINE concatMapAndUnzipM #-}
concatMapAndUnzipM :: Monad m => (a -> m ([b], [c])) -> [a] -> m ([b], [c])
concatMapAndUnzipM f = fold2M f (++) (++) ([], [])

{-# INLINE concatMapAndUnzip3M #-}
concatMapAndUnzip3M :: (Monad m) => (a -> m ([b], [c], [d])) -> [a] -> m ([b], [c], [d])    
concatMapAndUnzip3M f = fold3M f (++) (++) (++) ([], [], [])

{-# INLINE concatMapAndUnzip4M #-}
concatMapAndUnzip4M :: (Monad m) => (a -> m ([b], [c], [d], [e])) -> [a] -> m ([b], [c], [d], [e])    
concatMapAndUnzip4M f = fold4M f (++) (++) (++) (++) ([], [], [], [])

{-# INLINE concatMapAndUnzip3MD #-}
concatMapAndUnzip3MD :: (Monad m) => (a -> m (DL.DList b, DL.DList c, DL.DList d)) -> [a] -> m (DL.DList b, DL.DList c, DL.DList d)    
concatMapAndUnzip3MD f = fold3M f (DL.append) (DL.append) (DL.append) (DL.empty, DL.empty, DL.empty)

{-# INLINE concatMapAndUnzip3MD' #-}
concatMapAndUnzip3MD' :: (Monad m) => (a -> m (DL.DList b, DL.DList c, DL.DList d)) -> [a] -> m (DL.DList b, DL.DList c, DL.DList d)    
concatMapAndUnzip3MD' f = fold3M' f (DL.append) (DL.append) (DL.append) (DL.empty, DL.empty, DL.empty)

{-# INLINE concatMapAndUnzip4MD #-}
concatMapAndUnzip4MD :: (Monad m) => (a -> m (DL.DList b, DL.DList c, DL.DList d, DL.DList e)) -> [a] -> m (DL.DList b, DL.DList c, DL.DList d, DL.DList e)    
concatMapAndUnzip4MD f = fold4M f (DL.append) (DL.append) (DL.append) (DL.append) (DL.empty, DL.empty, DL.empty, DL.empty)

{-# INLINE concatMap2M #-}
concatMap2M :: Monad m => (a -> m (b, [c])) -> [a] -> m ([b], [c])
concatMap2M f = fold2M f (:) (++) ([], [])

{-# INLINE concatMap3M #-}
concatMap3M :: Monad m => (a -> m (b, [c], [d])) -> [a] -> m ([b], [c], [d])
concatMap3M f = fold3M f (:) (++) (++) ([], [], [])

crush3 :: [(a, [b], [c])] -> ([a], [b], [c])
crush3 lst = let 
         c1 = map (\ (x, _, _) -> x) lst
         c2 = concatMap (\ (_, x, _) -> x) lst
         c3 = concatMap (\ (_, _, x) -> x) lst
    in (c1, c2, c3)

crush3D :: [(DList a, DList b, DList c)] -> (DList a, DList b, DList c)
crush3D = foldl' worker (DL.empty, DL.empty, DL.empty)
    where
    worker (a, b, c) (dla, dlb, dlc) = 
        (dla `DL.append` a, dlb `DL.append` b, dlc `DL.append` c)

crush2 :: [(a, [b])] -> ([a], [b])
crush2 lst = let 
         c1 = map fst lst
         c2 = concatMap snd lst
    in (c1, c2)

zipWithAndUnzipM :: Monad m => (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
zipWithAndUnzipM f l1 l2 = liftM unzip $ zipWithM f l1 l2

zipWithAndUnzip3M :: Monad m => (a -> b -> m (c, d, e)) -> [a] -> [b] -> m ([c], [d], [e])
zipWithAndUnzip3M f l1 l2 = liftM unzip3 $ zipWithM f l1 l2

zipWithSeq :: (Integer -> a -> b) -> [a] -> [b]
zipWithSeq f = zipWith f [0..]

zipWithSeqM :: Monad m => (Integer -> a -> m b) -> [a] -> m [b]
zipWithSeqM f = zipWithM f [0..]

partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a],[a])
partitionM p = foldr worker (return ([], []))
    where
    worker x = liftM2 select (p x)
        where
        select b ~(ts,fs) = if b
            then (x:ts,fs)
            else (ts, x:fs)

foldMM :: Monad m => (a -> m b') -> (b' -> b -> b) -> b -> [a] -> m b
foldMM f op z = worker
    where
    worker []     = return z
    worker (x:xs) = liftM2 op (f x) (worker xs)

foldMM' :: Monad m => (a -> m b') -> (b' -> b -> b) -> b -> [a] -> m b
foldMM' f op z = foldM worker z
    where
    worker b a = liftM2 op (f a) (return b)

fold2M :: Monad m => (a -> m (r1', r2')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r1, r2) -> [a] -> m (r1, r2)
fold2M f op1 op2 = foldMM f op
    where
    op (r1', r2') (r1, r2) = (r1' `op1` r1, r2' `op2` r2)

fold2M' :: Monad m => (a -> m (r1', r2')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r1, r2) -> [a] -> m (r1, r2)
fold2M' f op1 op2 = foldMM' f op
    where
    op (r1', r2') (r1, r2) = (r1' `op1` r1, r2' `op2` r2)

fold3M :: Monad m => (a -> m (r1', r2', r3')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r3' -> r3 -> r3) 
                  -> (r1, r2, r3) -> [a] -> m (r1, r2, r3)
fold3M f op1 op2 op3 = foldMM f op
    where
    op (r1', r2', r3') (r1, r2, r3) = (r1' `op1` r1, r2' `op2` r2, r3' `op3` r3)

fold3M' :: Monad m => (a -> m (r1', r2', r3')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r3' -> r3 -> r3) 
                  -> (r1, r2, r3) -> [a] -> m (r1, r2, r3)
fold3M' f op1 op2 op3 = foldMM' f op
    where
    op (r1', r2', r3') (r1, r2, r3) = (r1' `op1` r1, r2' `op2` r2, r3' `op3` r3)

fold4M :: Monad m => (a -> m (r1', r2', r3', r4')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r3' -> r3 -> r3) 
                  -> (r4' -> r4 -> r4)
                  -> (r1, r2, r3, r4) -> [a] -> m (r1, r2, r3, r4)
fold4M f op1 op2 op3 op4 = foldMM f op
    where
    op (r1', r2', r3', r4') (r1, r2, r3, r4) = (r1' `op1` r1, r2' `op2` r2, r3' `op3` r3, r4' `op4` r4)

fold4M' :: Monad m => (a -> m (r1', r2', r3', r4')) 
                  -> (r1' -> r1 -> r1) 
                  -> (r2' -> r2 -> r2) 
                  -> (r3' -> r3 -> r3) 
                  -> (r4' -> r4 -> r4)
                  -> (r1, r2, r3, r4) -> [a] -> m (r1, r2, r3, r4)
fold4M' f op1 op2 op3 op4 = foldMM' f op
    where
    op (r1', r2', r3', r4') (r1, r2, r3, r4) = (r1' `op1` r1, r2' `op2` r2, r3' `op3` r3, r4' `op4` r4)

mapMaybeM :: Monad m => (a -> m b) -> Maybe a -> m (Maybe b)
mapMaybeM _ Nothing = return Nothing
mapMaybeM f (Just v) = liftM Just $ f v

joinJust :: Maybe a -> Maybe b -> Maybe (a, b)
joinJust (Just a) (Just b) = Just (a, b)
joinJust Nothing Nothing = Nothing
joinJust _ _ = error "joinJust: unexpected case"

consMaybe :: Maybe a -> [a] -> [a]
consMaybe Nothing  l = l
consMaybe (Just e) l = e : l

if' :: Bool -> (a -> b) -> (a -> b) -> a -> b
if' b f g x = if b then f x else g x

ifM :: (a -> Bool) -> (a -> b) -> (a -> b) -> a -> b
ifM bf f g x = if bf x then f x else g x

findAndDelete :: (a -> Bool) -> [a] -> (Maybe a, [a])
findAndDelete _ [] = (Nothing, [])
findAndDelete p (x:xs) = if p x then (Just x, xs) else let
        (r, xs') = findAndDelete p xs
    in (r, x : xs')

-- Returns the n-th element of a list
at :: Integer -> [a] -> a
at = flip genericIndex 

-- Returns the range between n-th and m-th elements of a list
range :: Integer -> Integer -> [a] -> [a]
range n m = genericDrop n . genericTake (m+1)

-- Finds the position of x in a list
getPos :: (Eq a) => a -> [a] -> Integer
getPos x = toInteger . fromMaybe (error "<Utilities>.<getPos>: Unexpected case") . elemIndex x 

-- Changes the n-th element of a list
putAt :: Integer -> a -> [a] -> [a]
putAt n x l = genericTake n l ++ [x] ++ genericDrop (n+1) l

-- XXX: check this
replaceAt :: Int -> a -> [a] -> [a]
replaceAt pos val lst = let 
        (pref, suff) = splitAt pos lst
    in pref ++ val : tail suff

-- Changes the n-th to m-th elements of a list
putRange :: Integer -> Integer -> [a] -> [a] -> [a]
putRange n m xs l = genericTake n l ++ xs ++ genericDrop (m+1) l

-- Split a list at regular intervals
chunk :: Integer -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = let (y, ys) = genericSplitAt n xs in y : chunk n ys


-- Auxiliary functions ---------------------------------------------------------

nestStr :: Int -> String -> String
nestStr n = unlines . map (replicate n ' ' ++) . lines