{-# LANGUAGE TypeFamilies, MultiParamTypeClasses, FlexibleInstances,
             IncoherentInstances, RankNTypes, ScopedTypeVariables,
             FlexibleContexts,UndecidableInstances, DeriveDataTypeable,
             ImpredicativeTypes #-}
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 Control.Monad.IO.Class
import Data.IORef()
import Data.Supply as S
import Data.Data
import Data.IORef
import Data.Dynamic

asType :: a -> a -> a
asType a _ = 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 = foldl App (Var var)

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 (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 =
        checkCond1 v (Condition d x y)
    checkCond1 v p@(Condition d x z@(App (Shift e f g) h)) =
        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 (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) | GEC a b `elem` xs            = True
    checkConds (GEC a b : xs) | LTC a b `elem` 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]

data LU = LU deriving (Eq, Ord, Data, Typeable)
instance SBV.SymWord LU
instance SBV.HasKind LU

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
    varmap <- liftIO createVarPool
    expmap <- liftIO createExpPool
    supply  <- liftIO $ newNumSupply
    let (s1,s2,s3,s4) = split4 supply
    cs <- mapM (\(s,x) -> compilec varmap expmap s x) $ zip (split s1) d
    mapM_ (SBV.constrain) cs
    lhs <- compilel varmap expmap s3 a
    rhs <- compilel varmap expmap s4 b
    return $ lhs .== rhs

data VarT = VarI | VarF VarT VarT
type VarPool = IORef (Map.Map Int Dynamic)
type ExpPool = IORef (Map.Map L SBV.SInteger)

createVarPool :: IO (VarPool)
createVarPool = newIORef $ Map.empty
createExpPool :: IO (ExpPool)
createExpPool = newIORef $ Map.empty


sbvTc :: TyCon
sbvTc = mkTyCon3 "Data" "SBV" "SBV"

instance (Typeable t) => Typeable (SBV.SBV t) where
  typeOf x = mkTyConApp sbvTc [typeOf (get x)]
    where
      get :: SBV.SBV a -> a
      get = undefined


getVarSymbol :: VarPool -> (SBV.Symbolic Dynamic) -> Int -> SBV.Symbolic Dynamic
getVarSymbol pool factory a = do
  v <- liftIO $ readIORef pool
  case Map.lookup a v of
    Just q -> return $ q
    Nothing -> do
      dx <- factory
      let newmap = Map.insert a dx v
      liftIO $ writeIORef pool newmap
      return dx

getExpSymbol :: ExpPool -> (SBV.Symbolic SBV.SInteger) -> L -> SBV.Symbolic SBV.SInteger
getExpSymbol pool factory a = do
  v <- liftIO $ readIORef pool
  case Map.lookup a v of
    Just q -> return $ q
    Nothing -> do
      dx <- factory
      let newmap = Map.insert a dx v
      liftIO $ writeIORef pool newmap
      return dx


compilec :: VarPool -> ExpPool -> Supply Int -> Constraint -> SBV.Symbolic (SBV.SBool)
compilec v e s (Zero a)  = do
  lhs <- compilel v e s a
  return $ lhs .== (0::SBV.SInteger)

compilec v e s (LTC a b) = do
  let (s1,s2) = split2 s 
  lhs <- compilel v e s1 a
  rhs <- compilel v e s2 b
  return $ lhs .< rhs

compilec v e s (GEC a b) = do
  let (s1,s2) = split2 s 
  lhs <- compilel v e s1 a
  rhs <- compilel v e s2 b
  return $ lhs .>= rhs

class Typeable a => VarFactory a where
  createDyn :: String -> a -> SBV.Symbolic Dynamic
  createDyn a n = do
    x <- createVar a n
    return $ toDyn x
  createVar :: String -> a -> SBV.Symbolic a

instance VarFactory SBV.SInteger where
  createVar n _ = SBV.free n

compilel :: VarPool -> ExpPool -> Supply Int -> L -> SBV.Symbolic (SBV.SInteger)
compilel v e s (Op a c b) = do
  let (s1, s2) = split2 s
  al <- compilel v e s1 a
  bl <- compilel v e s1 b
  return $ case c of
    '+' -> al + bl
    '-' -> al - bl
    '*' -> al * bl
    '/' -> al `SBV.sDiv` bl
compilel v e s (Var a)   = do
  let itype = (undefined ::SBV.SInteger)
  sym <- getVarSymbol v (createDyn (showVar a "") itype) a
  let var = (fromDynamic sym) :: Maybe SBV.SInteger
  case var of
      Just x -> return x
      Nothing -> do
        error "Type Error"
compilel v e s (Num a)   = return $ SBV.literal $ fromIntegral a
compilel v e s (Bottom)  = SBV.free_
compilel v e s l = do
    sym <- getExpSymbol e (return $ SBV.uninterpret ("unint" ++ showVar (supplyValue s) "") ) l
    return sym
--compilel v e x         = error $  "Cannot compile "++ show x

gettype :: ((b -> t) -> t) -> a -> b -> a
gettype = error "???"

{-
compileapp :: (VarFactory b, SBV.Uninterpreted (b -> SBV.SInteger)) => VarPool -> L -> ((b -> t) -> t) -> SBV.Symbolic (SBV.SInteger)
compileapp v (Var a) f = do
  let itype = (error :: SBV.SInteger)
  sym <- getVarSymbol v (createDyn (showVar a "") (gettype f itype)) a
  let var = (fromDynamic sym) :: Maybe SBV.SInteger
  case var of
      Just x -> return x
      Nothing -> do
        error "Type Error"
-}
compileapp v (App x y) t = error $ "Not yet implemented."