{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Algorithm.Cooper
-- Copyright   :  (c) Masahiro Sakai 2011
-- License     :  BSD-style
-- 
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable (FlexibleInstances)
--
-- Naive implementation of Cooper's algorithm
--
-- Reference:
-- 
-- * <http://hagi.is.s.u-tokyo.ac.jp/pub/staff/hagiya/kougiroku/ronri/5.txt>
-- 
-- * <http://www.cs.cmu.edu/~emc/spring06/home1_files/Presburger%20Arithmetic.ppt>
-- 
-- See also:
--
-- * <http://hackage.haskell.org/package/presburger>
--
-----------------------------------------------------------------------------
module Algorithm.Cooper
    (
    -- * Language of presburger arithmetics
      ExprZ
    , Lit (..)
    , QFFormula (..)
    , (.|.)
    , Model

    -- * Projection
    , project
    , projectCases
    , projectCases'
    , projectCasesN

    -- * Quantifier elimination
    , eliminateQuantifiers

    -- * Constraint solving
    , solveFormula
    , solveQFFormula
    , solveConj
    , solveQFLA
    ) where

import Control.Monad
import Data.List
import Data.Maybe
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Ratio

import Data.ArithRel
import Data.Expr
import Data.Formula
import Data.Linear
import qualified Data.LA as LA
import qualified Algorithm.FourierMotzkin as FM

-- ---------------------------------------------------------------------------

-- | Linear arithmetic expression over integers.
type ExprZ = LA.Expr Integer

atomZ :: RelOp -> Expr Rational -> Expr Rational -> Maybe QFFormula
atomZ op a b = do
  (e1,c1) <- FM.termR a
  (e2,c2) <- FM.termR b
  let a' = c2 .*. e1
      b' = c1 .*. e2
  case op of
    Le -> return $ Lit $ a' `leZ` b'
    Lt -> return $ Lit $ a' `ltZ` b'
    Ge -> return $ Lit $ a' `geZ` b'
    Gt -> return $ Lit $ a' `gtZ` b'
    Eql -> return $ eqZ a' b'
    NEq -> liftM notB (atomZ Eql a b)

leZ, ltZ, geZ, gtZ :: ExprZ -> ExprZ -> Lit
leZ e1 e2 = e1 `ltZ` (e2 .+. LA.constant 1)
ltZ e1 e2 = Pos $ (e2 .-. e1)
geZ = flip leZ
gtZ = flip gtZ

eqZ :: ExprZ -> ExprZ -> QFFormula
eqZ e1 e2 = Lit (e1 `leZ` e2) .&&. Lit (e1 `geZ` e2)

-- | Literal
-- 
-- * @Pos e@ means @e > 0@
-- 
-- * @Divisible True d e@ means @e@ can be divided by @d@ (i.e. @d|e@)
-- 
-- * @Divisible False d e@ means @e@ can not be divided by @d@ (i.e. @¬(d|e)@)
data Lit
    = Pos ExprZ
    | Divisible Bool Integer ExprZ
    deriving (Show, Eq, Ord)

instance Variables Lit where
  vars (Pos t) = vars t
  vars (Divisible _ _ t) = vars t

instance Complement Lit where
  notB (Pos e) = e `leZ` LA.constant 0
  notB (Divisible b c e) = Divisible (not b) c e

-- | quantifier-free negation normal form
data QFFormula
    = T'
    | F'
    | And' QFFormula QFFormula
    | Or' QFFormula QFFormula
    | Lit Lit
    deriving (Show, Eq, Ord)

instance Complement QFFormula where
  notB T' = F'
  notB F' = T'
  notB (And' a b) = Or' (notB a) (notB b)
  notB (Or' a b) = And' (notB a) (notB b)
  notB (Lit lit) = Lit (notB lit)

instance Lattice QFFormula where
  top    = T'
  bottom = F'
  meet   = And'
  join   = Or'

instance Boolean QFFormula

instance Variables QFFormula where
  vars T' = IS.empty
  vars F' = IS.empty
  vars (And' a b) = vars a `IS.union` vars b
  vars (Or' a b)  = vars a `IS.union` vars b
  vars (Lit l)    = vars l

instance IsRel (LA.Expr Integer) QFFormula where
  rel op lhs rhs =
    case op of
      Le  -> Lit $ leZ lhs rhs
      Ge  -> Lit $ geZ lhs rhs
      Lt  -> Lit $ ltZ lhs rhs
      Gt  -> Lit $ gtZ lhs rhs
      Eql -> eqZ lhs rhs
      NEq -> notB $ rel Eql lhs rhs

(.|.) :: Integer -> ExprZ -> QFFormula
n .|. e = Lit $ Divisible True n e

subst1 :: Var -> ExprZ -> QFFormula -> QFFormula
subst1 x e = go
  where
    go T' = T'
    go F' = F'
    go (And' a b) = And' (go a) (go b)
    go (Or' a b) = Or' (go a) (go b)
    go (Lit (Divisible b c e1)) = Lit $ Divisible b c $ LA.applySubst1 x e e1
    go (Lit (Pos e1)) = Lit $ Pos $ LA.applySubst1 x e e1

simplify :: QFFormula -> QFFormula
simplify (And' a b) = simplify1 $ And' (simplify a) (simplify b)
simplify (Or' a b)  = simplify1 $ Or' (simplify a) (simplify b)
simplify formula    = simplify1 formula

simplify1 :: QFFormula -> QFFormula
simplify1 T' = T'
simplify1 F' = F'
simplify1 (And' a b) =
  case (a, b) of
    (T', b') -> b'
    (a', T') -> a'
    (F', _) -> false
    (_, F') -> false
    (a',b') -> a' .&&. b'
simplify1 (Or' a b) =
  case (a, b) of
    (F', b') -> b'
    (a', F') -> a'
    (T', _) -> true
    (_, T') -> true
    (a',b') -> a' .||. b'
simplify1 (Lit lit) = simplifyLit lit

simplifyLit :: Lit -> QFFormula
simplifyLit (Pos e) =
  case LA.asConst e of
    Just c -> if c > 0 then true else false
    Nothing ->
      -- e > 0  <=>  e - 1 >= 0
      -- <=>  LA.mapCoeff (`div` d) (e - 1) >= 0
      -- <=>  LA.mapCoeff (`div` d) (e - 1) + 1 > 0
      Lit $ Pos $ LA.mapCoeff (`div` d) (e .-. LA.constant 1) .+. LA.constant 1
  where
    d = if null cs then 1 else abs $ foldl1' gcd cs
    cs = [c | (c,x) <- LA.terms e, x /= LA.unitVar]
simplifyLit lit@(Divisible b c e)
  | LA.coeff LA.unitVar e `mod` d /= 0 = if b then false else true
  | c' == 1   = if b then true else false
  | d  == 1   = Lit lit
  | otherwise = Lit $ Divisible b c' e'
  where
    d  = abs $ foldl' gcd c [c2 | (c2,x) <- LA.terms e, x /= LA.unitVar]
    c' = c `div` d
    e' = LA.mapCoeff (`div` d) e

-- ---------------------------------------------------------------------------

data Witness = WCase1 Integer ExprZ | WCase2 Integer Integer Integer [ExprZ]

evalWitness :: Model Integer -> Witness -> Integer
evalWitness model (WCase1 c e) = LA.evalExpr model e `div` c
evalWitness model (WCase2 c j delta us)
  | null us'  = j `div` c
  | otherwise = (j + (((u - delta - 1) `div` delta) * delta)) `div` c
  where
    us' = map (LA.evalExpr model) us
    u = minimum us'

-- ---------------------------------------------------------------------------

project :: Var -> QFFormula -> QFFormula
project x formula = simplify $ orB [phi | (phi,_) <- projectCases' x formula, phi /= F']

projectCases :: Var -> QFFormula -> [(QFFormula, Model Integer -> Model Integer)]
projectCases x formula = do
  (phi, wit) <- projectCases' x formula
  return (phi, \m -> IM.insert x (evalWitness m wit) m)

projectCases' :: Var -> QFFormula -> [(QFFormula, Witness)]
projectCases' x formula = [(simplify phi, w) | (phi,w) <- case1 ++ case2]
  where
    -- xの係数の最小公倍数
    c :: Integer
    c = f formula
      where
         f :: QFFormula -> Integer
         f T' = 1
         f F' = 1
         f (And' a b) = lcm (f a) (f b)
         f (Or' a b) = lcm (f a) (f b)
         f (Lit (Pos e)) = fromMaybe 1 (LA.lookupCoeff x e)
         f (Lit (Divisible _ _ e)) = fromMaybe 1 (LA.lookupCoeff x e)
    
    -- 式をスケールしてxの係数の絶対値をcへと変換し、その後cxをxで置き換え、
    -- xがcで割り切れるという制約を追加した論理式
    formula1 :: QFFormula
    formula1 = simplify $ f formula .&&. Lit (Divisible True c (LA.var x))
      where
        f :: QFFormula -> QFFormula
        f T' = T'
        f F' = F'
        f (And' a b) = f a .&&. f b
        f (Or' a b) = f a .||. f b
        f lit@(Lit (Pos e)) =
          case LA.lookupCoeff x e of
            Nothing -> lit
            Just a ->
              let s = abs (c `div` a)
              in Lit $ Pos $ g s e
        f lit@(Lit (Divisible b d e)) =
          case LA.lookupCoeff x e of
            Nothing -> lit
            Just a ->
              let s = abs (c `div` a)
              in Lit $ Divisible b (s*d) $ g s e

        g :: Integer -> ExprZ -> ExprZ
        g s = LA.mapCoeffWithVar (\c' x' -> if x==x' then signum c' else s*c')

    -- d|x+t という形の論理式の d の最小公倍数
    delta :: Integer
    delta = f formula1
      where
        f :: QFFormula -> Integer
        f T' = 1
        f F' = 1
        f (And' a b) = lcm (f a) (f b)
        f (Or' a b)  = lcm (f a) (f b)
        f (Lit (Divisible _ d _)) = d
        f (Lit (Pos _)) = 1

    -- ts = {t | t < x は formula1 に現れる原子論理式}
    ts :: [ExprZ]
    ts = f formula1
      where
        f :: QFFormula -> [ExprZ]
        f T' = []
        f F' = []
        f (And' a b) = f a ++ f b
        f (Or' a b) = f a ++ f b
        f (Lit (Divisible _ _ _)) = []
        f (Lit (Pos e)) =
          case LA.extractMaybe x e of
            Nothing -> []
            Just (1, e')  -> [lnegate e'] -- Pos e <=> (x + e' > 0) <=> (-e' < x)
            Just (-1, _) -> [] -- Pos e <=> (-x + e' > 0) <=> (x < e')
            _ -> error "should not happen"

    -- formula1を真にする最小のxが存在する場合
    case1 :: [(QFFormula, Witness)]
    case1 = [ (subst1 x e formula1, WCase1 c e)
            | j <- [1..delta], t <- ts, let e = t .+. LA.constant j ]

    -- formula1のなかの x < t を真に t < x を偽に置き換えた論理式
    formula2 :: QFFormula
    formula2 = simplify $ f formula1
      where        
        f :: QFFormula -> QFFormula
        f T' = T'
        f F' = F'
        f (And' a b) = f a .&&. f b
        f (Or' a b) = f a .||. f b
        f lit@(Lit (Pos e)) =
          case LA.lookupCoeff x e of
            Nothing -> lit
            Just 1    -> F' -- Pos e <=> ( x + e' > 0) <=> -e' < x
            Just (-1) -> T' -- Pos e <=> (-x + e' > 0) <=>  x  < e'
            _ -> error "should not happen"
        f lit@(Lit (Divisible _ _ _)) = lit

    -- us = {u | x < u は formula1 に現れる原子論理式}
    us :: [ExprZ]
    us = f formula1
      where
        f :: QFFormula -> [ExprZ]
        f T' = []
        f F' = []
        f (And' a b) = f a ++ f b
        f (Or' a b) = f a ++ f b
        f (Lit (Pos e)) =
          case LA.extractMaybe x e of
            Nothing -> []
            Just (1, _)   -> []   -- Pos e <=> ( x + e' > 0) <=> -e' < x
            Just (-1, e') -> [e'] -- Pos e <=> (-x + e' > 0) <=>  x  < e'
            _ -> error "should not happen"
        f (Lit (Divisible _ _ _)) = []

    -- formula1を真にする最小のxが存在しない場合
    case2 :: [(QFFormula, Witness)]
    case2 = [(subst1 x (LA.constant j) formula2, WCase2 c j delta us) | j <- [1..delta]]

projectCasesN :: VarSet -> QFFormula -> [(QFFormula, Model Integer -> Model Integer)]
projectCasesN vs2 = f (IS.toList vs2)
  where
    f :: [Var] -> QFFormula -> [(QFFormula, Model Integer -> Model Integer)]
    f [] formula = return (formula, id)
    f (v:vs) formula = do
      (formula2, mt1) <- projectCases v formula
      (formula3, mt2) <- f vs formula2
      return (formula3, mt1 . mt2)

-- ---------------------------------------------------------------------------

-- | eliminate quantifiers and returns equivalent quantifier-free formula.
eliminateQuantifiers :: Formula (Atom Rational) -> Maybe QFFormula
eliminateQuantifiers = f
  where
    f T = return T'
    f F = return F'
    f (Atom (Rel e1 op e2)) = atomZ op e1 e2
    f (And a b) = liftM2 (.&&.) (f a) (f b)
    f (Or a b) = liftM2 (.||.) (f a) (f b)
    f (Not a) = f (pushNot a)
    f (Imply a b) = f $ Or (Not a) b
    f (Equiv a b) = f $ And (Imply a b) (Imply b a)
    f (Forall x body) = liftM notB $ f $ Exists x $ Not body
    f (Exists x body) = liftM (project x) (f body)

-- ---------------------------------------------------------------------------

solveFormula :: Formula (Atom Rational) -> SatResult Integer
solveFormula formula =
  case eliminateQuantifiers formula of
    Nothing -> Unknown
    Just formula' ->
       case solveQFFormula formula' of
         Nothing -> Unsat
         Just m -> Sat m

solveQFFormula :: QFFormula -> Maybe (Model Integer)
solveQFFormula formula = listToMaybe $ do
  (formula2, mt) <- projectCasesN (vars formula) formula
  case formula2 of
    T' -> return $ mt IM.empty
    _  -> mzero

-- | solve a (open) quantifier-free formula
solveConj :: [LA.Atom Rational] -> Maybe (Model Integer)
solveConj cs = solveQFFormula formula
  where
    formula = andB [rel op (f (lhs .-. rhs)) (LA.constant 0) | Rel lhs op rhs <- cs]
    f e = LA.mapCoeff (round . (s*)) e
      where
        s = fromInteger $ foldl' lcm 1 [denominator c | (c,_) <- LA.terms e]

-- | solve a (open) quantifier-free formula
solveQFLA :: [LA.Atom Rational] -> VarSet -> Maybe (Model Rational)
solveQFLA cs ivs = listToMaybe $ do
  (cs2, mt) <- FM.projectN rvs cs
  m <- maybeToList $ solveConj cs2
  return $ mt $ IM.map fromInteger m
  where
    rvs = vars cs `IS.difference` ivs

-- ---------------------------------------------------------------------------

testHagiya :: QFFormula
testHagiya = project 1 $ andB [c1, c2, c3]
  where
    [x,y,z] = map LA.var [1..3]
    c1 = x .<. (y .+. y)
    c2 = z .<. x
    c3 = 3 .|. x

{-
∃ x. 0 < y+y ∧ z<x ∧ 3|x
⇔
(2y-z > 0 ∧ 3|z+1) ∨ (2y-z > -2 ∧ 3|z+2) ∨ (2y-z > -3 ∧ 3|z+3)
-}

test3 :: QFFormula
test3 = project 1 $ andB [p1,p2,p3,p4]
  where
    x = LA.var 0
    y = LA.var 1
    p1 = LA.constant 0 .<. y
    p2 = 2 .*. x .<. y
    p3 = y .<. x .+. LA.constant 2
    p4 = 2 .|. y

{-
∃ y. 0 < y ∧ 2x<y ∧ y < x+2 ∧ 2|y
⇔
(2x < 2 ∧ 0 < x) ∨ (0 < 2x+2 ∧ x < 0)
-}

-- ---------------------------------------------------------------------------