{-# LANGUAGE ScopedTypeVariables, ExistentialQuantification, Rank2Types #-}
module Examples where

import Lambda
import SizedExp
import Constraints()
import Prelude ( ($), (+), (-), Int, (==), (*), (<), (>), (<=), (>=), (/=) )
import qualified Prelude as P
import qualified Control.Monad as M
import qualified Data.List as List

head :: (SizedExp se ) => Size se ([l] -> l)
head = bind headc body
    where 
    body l = match l true P.const

tail :: (SizedExp se ) => Size se ([l] -> [l])
tail = bind tailc body
    where body l = match l true (\_ xs -> xs)

cons :: (SizedExp se) => Size se (x -> [x] -> [x])
cons = bind conss true

t3 :: (SizedExp se) => Size se ( (a -> a) -> a -> a )
t3 = bind t3s body
    where body f x = f `app` (f `app` (f `app` x))

nil :: (SizedExp se) => Size se [x]
nil  = bind nils true

map :: (SizedExp se)  => Size se ( (a->b) -> [a] -> [b] )
map = bind smap body
    where body f l = match l nil
            ( \x xs ->  cons `app` (f `app` x) `app` (map `app` f `app` xs ))

heads :: (SizedExp se) => Size se ( [[a]] -> [a] )
heads = bind sheads $ \l ->  map `app` head `app` l

tails :: (SizedExp se) => Size se ( [[a]] -> [[a]] )
tails = bind stails $ \l ->  map `app` tail `app` l

append :: (SizedExp se) => Size se ([a] -> [a] -> [a])
append = bind appends body
    where body l1 l2 = match l1 l2
            (\x xs -> cons `app` x `app` (append `app` xs `app` l2))

t27 :: (SizedExp se) => Size se ((a -> a) -> a -> a)
t27 = bind (App t3s t3s) $ t3 `app` t3

t27_ :: (SizedExp se) => Size se ((a -> a) -> a -> a)
t27_ = bind (App t3s t3s) $ 
        \f ->  t3 `app` t3 `app` f

t27__ :: (SizedExp se) => Size se ((a -> a) -> a -> a)
t27__ = bind (App t3s t3s) $ \f x ->
            t3 `app` t3 `app` f `app` x

pam :: (SizedExp se) => Size se ([a -> b] -> a -> [b])
pam = bind (AAbs 1 2 $ Abs 3 $ List (Var 1) (Abs 4 $ App (App (Var 2) (Var 4)) (Var 3))) $ \fl x -> match fl nil
    (\f fs -> cons `app` (f `app` x) `app` (pam `app` fs `app` x))

reverse :: (SizedExp se) => Size se([a] -> [a])
reverse = bind reverses $ \l -> match l
                nil
            (
                \x xs -> append `app` (reverse `app` xs) `app` (cons `app` x `app` nil)
            )

addone :: (SizedExp se) => Size se ([P.Int] -> [P.Int])
addone = bind addones $ \l -> cons `app` 1 `app` l

add3 :: (SizedExp se) => Size se ([P.Int] -> [P.Int])
add3 = bind (AAbs 0 1 $ List (Op (Var 0) '+' (Num 3)) (Var 1)) $
        \l -> t3 `app` addone `app` l

add3_ :: (SizedExp se) => Size se ([P.Int] -> [P.Int])
add3_ = bind (AAbs 0 1 $ List (Op (Var 0) '+' (Num 3)) (Var 1)) $
        t3 `app` addone 

t9_ :: (SizedExp se) => Size se ((a -> a) -> a -> a)
t9_ = bind ( Abs 0 $ App t3s (App t3s (Var 0))) $
        \f -> t3 `app` (t3 `app` f)

t9 :: (SizedExp se) => Size se ((a -> a) -> a -> a)
t9 = bind ( Abs 0 $ App t3s (App t3s (Var 0))) $
        \f x -> t3 `app` (t3 `app` f) `app` x

add27s :: L
add27s = AAbs 0 1 $ List (Op (Var 0) '+' (Num 27)) (Abs 2 Unsized)
add27 :: (SizedExp se) => Size se ([P.Int] -> [P.Int])
add27 = bind add27s $ \x ->  t27 `app` addone `app` x


zipWiths :: L
zipWiths = let q = App (Var 4) $ Op (Op (Var 5) '+' (Var 3)) '-' (Var 1)  in
    (Abs 0 $ AAbs 1 2 $ AAbs 3 4 $ List (Var 1) (Abs 5 $ App (App (Var 0) (App (Var 2) (Var 5))) q ))
zipWith :: (SizedExp se) => Size se ((a2 -> a1 -> a) -> [a2] -> [a1] -> [a])
zipWith = bind zipWiths $ \f l1 l2 ->
        match l1
            nil
        (
            \x xs -> match l2 
                true
            (
                \y ys -> cons `app` (f `app` x `app` y) `app` (zipWith `app` f `app` xs `app` ys)
            )
        )
appAll :: (SizedExp se) => Size se ( [a -> b] -> a -> [b] )
appAll = bind (AAbs 0 1 $ Abs 2 $ List (Var 0) (Abs 3 $ Var 1 `App` Var 3 `App` Var 2 ) ) $ \fl x -> match fl 
            nil
        (
            \f fs -> cons `app` (f `app` x) `app` (appAll `app` fs `app` x)
        )

conspack :: (SizedExp se) => Size se (P.Int -> [P.Int] -> [P.Int])
conspack = bind (Abs 0 $ AAbs 1 2 $ List (Op (Var 1) '+' (Num 1)) (Abs 3 Unsized)) $ \x l ->
    match l (cons `app` x `app` l) 
    (\hd tl -> iff (x == hd) (
            cons `app` x `app` l
        ) (
            cons `app` hd `app` (conspack `app` x `app` tl)
        )
    )

cprod :: (SizedExp se) => Size se ([P.Int] -> [P.Int] -> [[P.Int]])
cprod = known  (AAbs 0 1 $ AAbs 2 3 $ List (Op (Var 0) '*' (Var 2)) $ Abs 4 $ List (Num 2) $ Abs 5 Unsized) 

sqdiff :: SizedExp se => Size se ([Int] -> [Int] -> [[Int]])
sqdiff = bind (let sq l = Op l '*' l in AAbs 0 1 $ AAbs 2 3 $ List (sq $ Op (Var 0) '-' (Var 2)) $ Abs 4 $ List (Num 2) $ Abs 5 Unsized) $
    \l1 l2 -> match l1 (cprod `app` l2 `app` l2)
        (\_ tl1 -> match l2 (cprod `app` l1 `app` l1)
            (\_ tl2 -> sqdiff `app` tl1 `app` tl2))

replace :: SizedExp se => Size se (Int -> [Int] -> [Int])
replace = bind (Abs 0 $ AAbs 1 2 $ List (Var 1) (Abs 3 Unsized)) $
    \x l -> match l nil (\hd tl -> cons `app` (x+hd) `app` tl)

scalarProd :: (SizedExp se0) =>  Size se0 ([Int] -> [Int] -> [Int])
scalarProd = bind (AAbs 0 1 $ AAbs 2 3 $ List (Num 1) (Abs 4  Unsized)) $ 
    \l1 l2 -> match l1 (
        match l2 ( cons `app` 0 `app` nil ) 
            (\_ _ -> true)
    ) ( \hd1 tl1 ->
        match l2 true
            ( \hd2 tl2 -> replace `app` (hd1 * hd2) `app` (scalarProd `app` tl1 `app` tl2) )
    )

mlist :: SizedExp se => Size se (a -> [a -> x] -> [x])
mlist = bind (Abs 0 $ AAbs 1 2 $ List (Var 1) (Abs 3 $ Var 2 `App` Var 3 `App` Var 0)) 
    $ \x l -> match l nil (\f fs -> cons `app` (f `app` x) `app` (mlist `app` x `app` fs))

strange :: (SizedExp se) => Size se ([Int] -> [Int])
strange = bind (AAbs 0 1 $ List (Num 2) (Abs 2 Unsized)) $ \l1 ->
    let b = match l1 nil (\x1 l2 -> match l2 nil (\x2 l3 ->  cons `app` x2 `app` (cons `app` x1 `app` nil)))
    in match b (cons `app` 0 `app` (cons `app` 0 `app` nil)) (\x xs -> b)


take4 :: SizedExp se => Size se (([a] -> [a]) -> [[a]])
take4 = bind (Abs 0 $ List (Num 1) (Abs 2 (Var 0 `App` (Var 0 `App` List (Num 0) (Abs 1 Bottom))))) $
    \f ->cons `app` (f `app` (f `app` nil) ) `app` nil


merge :: SizedExp se => Size se ([Int] -> [Int] -> [Int])
merge = bind (AAbs 0 1 $ AAbs 2 3 $ List (Op (Var 0) '+' (Var 2)) (Abs 4 Unsized))$
    \l1 l2 -> match l1 l2 (
        \x xs -> match l2 l1 (
                \y ys -> iff (x>y) (cons `app` x `app` (merge `app` xs `app` l2))
                                  (cons `app` y `app` (merge `app` l1 `app` ys))
            )
        )

split1 :: SizedExp se => Size se ([Int] -> [Int])
split1 = bind (AAbs 0 1 $ List (Op (Op (Var 0) '+' (Num 1))'/' (Num 2)) (Abs 2 Unsized))  $
    \z -> match z nil (\y ys -> cons `app` y `app` (split2 `app` ys))

split2 :: SizedExp se => Size se ([Int] -> [Int])
split2 = bind (AAbs 0 1 $ List (Op (Var 0) '/' (Num 2)) (Abs 2 Unsized)) $
    \z -> match z nil (\y ys -> split1 `app` ys)

ms = AAbs 0 1 $ List (Var 0) (Abs 4 Unsized)
mergesort :: SizedExp se => Size se ([Int] -> [Int])
mergesort = bind ms $
    \l -> match l nil (\x xs ->
        merge `app` (mergesort `app` (split1 `app` l)) `app` (mergesort `app` (split2 `app` l))
    )

last :: SizedExp se => Size se ([a] -> a)
last = bind (AAbs 0 1 $ App (Var 1) (Num 0)) $
    \l -> match l true (\x xs -> match xs x (\_ _ -> last `app` xs))

charm :: SizedExp se => Size se (([a] -> [a]) -> a)
charm = bind (Abs 0 $ App (AAbs 2 3 $ App (Var 3) (Num 0)) (App (Var 0) (List (Num 0) $ Abs 1 Bottom))) $
    \ f -> last `app` (f `app` nil)

fix :: SizedExp se => Size se ((a -> a) -> a)
fix = bind (Abs 2 $ App yComb (Var 2)) $ 
    \f -> f `app` (fix `app` f)

transpose :: SizedExp se => Size se ([[a]] -> [[a]])
transpose = bind transposec $ \l -> match l true $
    \l1 xss -> match l1 true $
        \x xs -> cons `app` (cons `app` x `app` (heads `app` xss))
            `app` (transpose `app` (cons `app` xs `app` (tails `app` xss)))
transposec = AAbs 18 5 $ List len fun
    where
    len = AAbs 19 6 (Var 19) `App` (Var 5 `App` Num 0)
    fun = Abs 8 $ List (Var 18) (Abs 9 $ AAbs 19 6 (Var 6 `App` Var 8) `App` (Var 5 `App` Var 9))

comps = Abs 2 $ Abs 3 $ Abs 4 $ App (Var 2) (App (Var 3) (Var 4))
comp :: (SizedExp se)  => Size se ( (b->c) -> (a->b) -> a->c )
comp = bind comps $ \f g x -> f `app` (g `app` x)

test1s = Abs 2 $ AAbs 19 6 $ Var 2 `App` List (Var 19) (Var 6)
test1 :: (SizedExp se)  => Size se (([a] -> [b]) -> [a] -> [b])
test1 = bind test1s $ \f l -> match (f `app` l) nil (\x xs -> f `app` l)


test2s = Abs 2 $ AAbs 18 5  $ appends `App` (Var 2 `App` List (Var 18) (Var 5)) `App`
  (appends `App` (Var 2 `App` List (Var 18) (Var 5)) `App` List (Var 18) (Var 5))
test2 :: (SizedExp se)  => Size se (([a] -> [a]) -> [a] -> [a])
test2 = bind test2s $ \f l -> append `app` (f `app` l) `app` (append `app` (f `app` l) `app` l)

data TestCase = forall a . TestCase P.String (forall se. SizedExp se => Size se a)

tests :: [TestCase]
tests = [
          TestCase "append" append 
        , TestCase "reverse" reverse 
        , TestCase "heads" heads 
        , TestCase "map" map 
        , TestCase "pam" pam
        , TestCase "head" head
        , TestCase "tail" tail
        , TestCase "t3" t3
        , TestCase "t9" t9
--        , TestCase "t9_" t9_  -- too few arguments in definition
--        , TestCase "t27" t27
--        , TestCase "t27_" t27_
        , TestCase "t27__" t27__
        , TestCase "addone" addone
        , TestCase "add3" add3
--        , TestCase "add3_" add3_ -- too few arguments in definition
        , TestCase "add27" add27
        , TestCase "zipWith" zipWith
        , TestCase "appAll" appAll
        , TestCase "conspack" conspack
        , TestCase "scalarProd" scalarProd
        , TestCase "sqdiff" sqdiff
        , TestCase "mlist" mlist
        , TestCase "strange" strange
        , TestCase "take4" take4
        , TestCase "charm" charm
        , TestCase "comp" comp
        , TestCase "merge" merge
        , TestCase "split1" split1
        , TestCase "split2" split2
        , TestCase "mergesort" mergesort
    ]

runTests = do
    failed <- M.forM tests $ \(TestCase name test) -> do
        P.print " +++++++++++++++++++++++++++++++"
        P.print $ " +  Proving " P.++ name
        P.print " +++++++++++++++++++++++++++++++"

        s <- prove test
        M.return [name | P.not s]

    let f = P.concat failed
    if List.null f then 
        P.putStrLn "All ok."
      else do
        P.putStr "\n\nFailed test cases: "
        P.putStrLn $ List.intercalate ", " f