{-# LANGUAGE TypeFamilies, MultiParamTypeClasses, FlexibleInstances,
             IncoherentInstances, RankNTypes, ScopedTypeVariables,
             FlexibleContexts,UndecidableInstances #-}
module Constraints where

import Lambda
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Map as Map
import qualified Data.Ord as Ord
import qualified Data.SBV as SBV
import Data.SBV ( (.==), (.<), (.>=))
import Control.Monad
import Data.IORef
import Data.Supply as S

asType :: a -> a -> a
asType a b = a

data D = D { var::Int, cond::[Condition] } 

data TypeKind = V | L TypeKind | F TypeKind TypeKind | U
    deriving Show


class TS a where
    tk :: a -> TypeKind
instance (TS a, TS b) => TS (a -> b) where
    tk _ = F (tk (undefined :: a)) (tk (undefined :: b))
instance TS a => TS [a] where
    tk _ = L (tk (undefined :: a))
instance TS a where
    tk _ = V
instance TS Int where
    tk _ = U

appall :: Int -> [L] -> L
appall var l = foldl (App) (Var $ var) $  l


data Condition = Condition [Constraint] L L

instance Show Condition where
    showsPrec _ (Condition d a b) = shows d . showString "|- " . shows a . showString " = " . shows b
    showList [] = showString ""
    showList [x] = shows x
    showList (x:xs) = shows x . showString "\n" . shows xs

data Constraint = Zero L | LTC L L | GEC L L
    deriving Eq

instance Show Constraint where
    showList []     = id
    showList [s]    = shows s . showChar ' '
    showList (x:xs) = shows x . showString ", " . shows xs
    showsPrec _ (Zero s) = shows s . showString " = 0"
    showsPrec _ (LTC s1 s2) = shows s1 . showString " < " . shows s2
    showsPrec _ (GEC s1 s2) = shows s1 . showString " >= " . shows s2

normalize :: L -> L
normalize l = delzero $ foldl (\a b -> Op a '-' b) (foldl (\a b -> Op  a '+' b) (f1 $ List.sortBy (Ord.comparing snd) a) c) d
    where
    (a,b,c,d) = norm l
    f1 ((0,l):xs) = f1 xs
    f1 ((1,l):xs) = Op (f1 xs) '+' (Var l)
    f1 ((-1,l):xs) = Op (f1 xs) '-' (Var l)
    f1 [] = Num b

    delzero (Op (Num 0) '+' b) = b
    delzero (Op a c b) = Op (delzero a) c b
    delzero l          = l

gnormalize :: L -> (L, Int)
gnormalize l = case (c,d) of
        ([], []) -> case nonzero of
            [] -> (Num 0, b)
            _  -> (expr, b)
        (_, _)   -> error $ "Expression in condition" ++ (show c)
    where
    nonzero = List.sortBy (Ord.comparing snd) $ filter ((/=0).fst) a
    proc (1, l) = Var l
    proc (-1, l) = Op (Num (-1)) '*' (Var l)
    t :: L -> (Int, Int) -> L
    t l (1, var) = Op l '+' (Var var)
    t l (-1, var) = Op l '-' (Var var)
    expr = foldl t (proc $ head nonzero) (tail nonzero)

    (a,b,c,d) = norm l

normalizec :: Constraint -> Constraint
normalizec (Zero l) = Zero $ normalize l
normalizec (LTC a b) = LTC (normalize a) (normalize b)
normalizec (GEC a b) = GEC (normalize a) (normalize b)
--normalizec (LTC a b) = LTC x (Num (-y))
--    where (x,y) = gnormalize $ Op a '-' b 
--normalizec (GEC a b) = GEC x (Num (-y))
--    where (x,y) = gnormalize $ Op a '-' b 

normalizecs :: [Constraint] -> [Constraint]
normalizecs = normalizecs' . Prelude.map normalizec
    where
    normalizecs' [] = []
    normalizecs' (x:xs) = x:normalizecs' (filter (/=x) xs)

{-
 - Takes an extended lambda expression and 
 -}
norm :: L -> ([(Int, Int)], Int, [L], [L])
norm (App a b)     = ([], 0, [App (normalize a) (normalize b)], [])
norm (List a b)    = ([], 0, [List (normalize a) (normalize b)], [])
norm (AAbs a b e)  = ([], 0, [AAbs a b (normalize e)], [])
norm (Shift a b c) = ([], 0, [Shift (normalize a) (normalize b) (normalize c)], [])
norm (Aggr a b c)  = ([], 0, [Aggr a (normalize b) (normalize c)], [])
norm (Unsized)     = ([], 0, [Unsized], [])
norm (Bottom)      = ([], 0, [Bottom], [])
norm (Abs i l)     = ([], 0, [Abs i $ normalize l], [])
norm (Var i)       = ([(1, i)], 0, [], [])
norm (Num i)       = ([], i, [], [])
norm q@(Op a c b)  = case c of
    '+' -> (a1++b1, a2+b2, a3++b3, a4++b4)
    '-' -> (sub a1 b1, a2-b2, a3++b4, a4++b3)
    '*' -> case a of
            (Num cnt) -> (map (mul cnt) b1, cnt*b2, map (Op (Num cnt) '*') b3, map (Op (Num cnt) '*') b4)
            _         -> ([], 0, [q], [])
    _   -> ([], 0, [q], [])
    where
        mul c (x,var) = (c*x, var)
        (a1, a2, a3, a4) = norm a
        (b1, b2, b3, b4) = norm b
        sub l1 ((c,i):xs) = sub (sub' l1 c i) xs
            where
            sub' [] c i = [(-c, i)]
            sub' ((cc,ii):xs) c i = if i==ii then (cc-c,i):xs else (cc,ii):sub' xs c i
        sub l1 [] = l1

tnorm = Op (Num 1) '+' (Op (Var 0) '-' (Num 1))

checkCond :: Supply Int -> [Condition] -> IO [Condition]
checkCond v l = do
    ll <- mapM (checkCond1 v) l
    return $ concat ll

    where
--    checkCond1 v p@(Condition d a b) | a==b = do
--        return  []
    checkCond1 v z@(Condition d (List a b) (List p q)) = do
--        i <- freshtypevar v
        let (v1,v2,v3) = split3 v
        let i = supplyValue v1
        let b' = rall $ App b (Var i)
        let q' = rall $ App q (Var i)
        l1 <- checkCond1 v2 (Condition ((Var i `GEC` Num 0):d) a p) 
        l2 <- checkCond1 v3 (Condition ((Var i `GEC` Num 0):(Var i `LTC` a):(Var i `LTC` p):d) b' q')
        return $ l1++l2
    checkCond1 v z@(Condition d (App (Shift e f g) h) x) = do
        let (v1,v2,v3) = split3 v
        let e' = rall $ App e h
        let g' = rall $ App g $ Op h '-' f
        l1 <- checkCond1 v1 (Condition ((h `LTC` f):d) e' x)
        l2 <- checkCond1 v2 (Condition ((h `GEC` f):d) g' x)
        putStrLn $ " -> " ++ (show l1)
        putStrLn $ " -> " ++ (show l2)
        return $ l1 ++ l2
    checkCond1 v z@(Condition d (App (Var a) x) (App (Var b) y)) | a==b = do
        checkCond1 v (Condition d x y)
    checkCond1 v p@(Condition d x z@(App (Shift e f g) h)) = do
        checkCond1 v (Condition d z x)
    checkCond1 v z@(Condition d a b) = return $ [Condition dd (normalize a) (normalize b)]
        where
        dd = normalizecs d

subst ndl hst (App a b)     = App (subst ndl hst a) (subst ndl hst b)
subst ndl hst (List a b)    = List (subst ndl hst a) (subst ndl hst b)
subst ndl hst (AAbs a b e)  = AAbs a b (subst ndl hst e)
subst ndl hst (Shift a b c) = Shift (subst ndl hst a) (subst ndl hst b) (subst ndl hst c)
subst ndl hst (Aggr a b c)  = Aggr a (subst ndl hst b) (subst ndl hst c)
subst ndl hst (Unsized)     = Unsized
subst ndl hst (Bottom)      = Bottom
subst ndl hst q@(Abs i l)   = if ndl==i then q else Abs i $ subst ndl hst l
subst ndl hst q@(Var i)     = if ndl==i then hst else q
subst ndl hst (Num i)       = Num i
subst ndl hst (Op a c b)    = Op (subst ndl hst a) c (subst ndl hst b)

substc ndl hst (Zero l)  = Zero $ subst ndl hst l
substc ndl hst (LTC a b) = LTC (subst ndl hst a) (subst ndl hst b)
substc ndl hst (GEC a b) = GEC (subst ndl hst a) (subst ndl hst b)

reorder cs = if any check cs 
        then Just (map r cs)
        else Nothing
    where
    check (Zero _)        = False
    check (LTC a (Num 0)) = False 
    check (GEC a (Num 0)) = False 
    check _               = True
    r (LTC a b) = LTC (normalize (Op a '-' b)) (Num 0)
    r (GEC a b) = GEC (normalize (Op a '-' b)) (Num 0)
    r l         = l

solve :: [Condition] -> Supply Int -> IO [Condition]
solve l supply = do
    ll <- forM (zip l $ split supply) (\(c,s) -> do
        putStrLn $ "\nSOLVING "++(show c)
        solve1 s c
        )
    return $ concat ll
    where

    searchzero (Zero (Var a):xs) = Just (a,xs)
    searchzero (x:xs)            = do
        (a,l) <- searchzero xs
        return (a,x:l)
    searchzero []                = Nothing


    searcheq (q@(GEC (Var var) exp):xs) prev = case findeq (List.reverse prev ++ xs) [] of
            Nothing -> searcheq xs (q:prev)
            l       -> l
        where
        expinc = normalize $ Op exp '+' $ Num 1
        findeq []                       _     = Nothing
        findeq (LTC (Var var2) exp2:xs) prev2 | var==var2 && exp2==expinc = Just (var, exp, List.reverse prev2 ++ xs)
        findeq (x:xs)                   prev2 = findeq xs (x:prev2)

    searcheq (x:xs)                 prev = searcheq xs (x:prev) 
    searcheq []                     _    = Nothing

    checkConds [] = False
    checkConds ((LTC a b):xs) | (elem (GEC a b) xs)            = True
    checkConds ((GEC a b):xs) | (elem (LTC a b) xs)            = True
    checkConds ((LTC a (Num b)):xs) | b<=0 && elem (Zero a) xs = True
    checkConds ((GEC a (Num b)):xs) | b>0  && elem (Zero a) xs = True
    checkConds (_:xs) = checkConds xs

    checkConds2 (LTC (Num a) (Num b):xs) | a>=b = Nothing
    checkConds2 (LTC (Num a) (Num b):xs) | a<b  = checkConds2 xs
    checkConds2 (GEC (Num a) (Num b):xs) | a<b  = Nothing
    checkConds2 (GEC (Num a) (Num b):xs) | a>=b = checkConds2 xs
    checkConds2 (x:xs)                          = do { y <- checkConds2 xs; return (x:y) }
    checkConds2 []                              = Just []

--    applyList a b d supp = do
--        let (s1,s2,s3) = split3 supp
--        let t = fresh (L V) s1
--        let dd = (Condition d (rall $ App a t) (rall $ App b t))
--        putStrLn $ "Applying a fresh variable: "++(show t) ++"\n"++(show dd)
--        x <- checkCond s2 [dd] >>= mapM (solve1 s3)
--        return $ concat x

    solve1  supp c@(Condition d a b) = case checkConds2 d of
        Nothing -> do
            putStrLn "Contradiction in preconditions"
            return []
        Just d' -> solve1' supp $ Condition d a b

    solve1' supp c@(Condition d a b)
        | checkConds d = do
            putStrLn "Contradiction in preconditions"
            return []
        | a==b = do
            putStrLn $ "Equals"
            return []
        | Just (var, nl) <- searchzero d = do
            let x = (Condition (Prelude.map (normalizec.substc var (Num 0)) nl)
                        (normalize$subst var (Num 0) a)
                        (normalize$subst var (Num 0) b)
                    )
            putStrLn $ (show $ Var var) ++" is zero:\n" ++ (show x)
            solve1 supp x 
        | Just (var, exp, nl) <- searcheq d [] = do
            putStrLn $ "Found equation "++(show $ Var var) ++ " = "++(show exp)
            let x = (Condition (Prelude.map (normalizec.substc var exp) nl )
                        (normalize$subst var exp a)
                        (normalize$subst var exp b)
                    )
            putStrLn $ "New equations:\n" ++ (show x)
            solve1 supp x 
        | App p q <- a, App r s <- b = do
            putStrLn $ "Branching!"
            nc <- checkCond supp [Condition d p r, Condition d q s] 
            solve nc supp


--        | Abs _ _ <- a, Abs _ _ <- b = do
--            let (s1, s2) = split2 supp
--            let t = fresh (tk a) s1
--            let dd= (Condition d (rall $ App a t) (rall $ App b t))
--            putStrLn $ "Applying a fresh variable: "++(show t) ++"\n"++(show dd)
--            solve1 s2 dd
--        | AAbs _ _ _ <- a, Abs  _ _   <- b = applyList a b d supp
--        | Abs  _ _   <- a, AAbs _ _ _ <- b = applyList a b d supp
--        | AAbs _ _ _ <- a, AAbs _ _ _ <- b = applyList a b d supp
        | Just dd <- reorder d = do 
            putStrLn $ "Reorder " ++ (show dd)
            solve1 supp $ Condition dd a b
        | otherwise = do
            putStrLn "Tying to call solver"
            let x = compiletosolver a b d
            y <- SBV.prove x
            print y
            case y of
                (SBV.ThmResult (SBV.Unsatisfiable _)) -> return []
                otherwise -> return [c]

fvc (Zero a)  = fv a
fvc (LTC a b) = fv a `Set.union` fv b
fvc (GEC a b) = fv a `Set.union` fv b

compiletosolver :: L -> L -> [Constraint] -> SBV.Symbolic SBV.SBool
compiletosolver a b d = do
    let fvs = Set.toList $ Set.unions $ (fv a):(fv b):(map fvc d)
    (vars::[SBV.SInteger]) <- mapM SBV.free $ map (flip showVar "") fvs
    let varmap = Map.fromList $ zip fvs vars
    mapM SBV.constrain $ map (compilec varmap) d
    return $ (compilel varmap a) .== (compilel varmap b)

compilec v (Zero a)  = compilel v a .== (0::SBV.SInteger)
compilec v (LTC a b) = compilel v a .<  compilel v b
compilec v (GEC a b) = compilel v a .>= compilel v b

compilel :: Map.Map Int SBV.SInteger -> L -> SBV.SInteger
compilel v (Op a c b) = case c of 
    '+' -> al + bl
    '-' -> al - bl
    '*' -> al * bl
    where
    al = compilel v a
    bl = compilel v b
compilel v (Var a) | Just x<-Map.lookup a v = x
compilel v (Num a) = SBV.literal $ fromIntegral a
compilel v a = error $ "Cannot compile "++(show a)