{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses #-}
{-# OPTIONS_GHC -Wall #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Arith.FourierMotzkin.Base
-- Copyright   :  (c) Masahiro Sakai 2011-2013
-- License     :  BSD-style
-- 
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable (TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses)
--
-- Naïve implementation of Fourier-Motzkin Variable Elimination
-- 
-- Reference:
--
-- * <http://users.cecs.anu.edu.au/~michaeln/pubs/arithmetic-dps.pdf>
--
-----------------------------------------------------------------------------
module ToySolver.Arith.FourierMotzkin.Base
    (
    -- * Primitive constraints
      Constr (..)
    , eqR
    , ExprZ
    , fromLAAtom
    , toLAAtom
    , constraintsToDNF
    , simplify

    -- * Bounds
    , Bounds
    , evalBounds
    , boundsToConstrs
    , collectBounds

    -- * Projection
    , project
    , project'
    , projectN
    , projectN'

    -- * Solving
    , solve
    , solve'

    -- * Utilities used by other modules
    , Rat
    , toRat
    ) where

import Control.Monad
import Data.List
import Data.Maybe
import Data.Ratio
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.VectorSpace hiding (project)
import qualified Data.Interval as Interval
import Data.Interval (Interval, Extended (..), (<=..<), (<..<=), (<..<))

import ToySolver.Data.OrdRel
import ToySolver.Data.Boolean
import ToySolver.Data.DNF
import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar

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

type ExprZ = LA.Expr Integer

-- | (t,c) represents t/c, and c must be >0.
type Rat = (ExprZ, Integer)

toRat :: LA.Expr Rational -> Rat
toRat e = seq m $ (LA.mapCoeff f e, m)
  where
    f x = numerator (fromInteger m * x)
    m = foldl' lcm 1 [denominator c | (c,_) <- LA.terms e]

fromRat :: Rat -> LA.Expr Rational
fromRat (e,c) = LA.mapCoeff (% c) e

evalRat :: Model Rational -> Rat -> Rational
evalRat model (e, d) = LA.lift1 1 (model IM.!) (LA.mapCoeff fromIntegral e) / (fromIntegral d)

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

-- | Atomic constraint
data Constr
  = IsNonneg ExprZ
  -- ^ e ≥ 0
  | IsPos ExprZ
  -- ^ e > 0
  | IsZero ExprZ
  -- ^ e = 0
  deriving (Show, Eq, Ord)

instance Variables Constr where
  vars (IsPos t) = vars t
  vars (IsNonneg t) = vars t
  vars (IsZero t) = vars t

leR, ltR, geR, gtR, eqR :: Rat -> Rat -> Constr
leR (e1,c) (e2,d) = IsNonneg $ normalizeExpr $ c *^ e2 ^-^ d *^ e1
ltR (e1,c) (e2,d) = IsPos $ normalizeExpr $ c *^ e2 ^-^ d *^ e1
geR = flip leR
gtR = flip ltR
eqR (e1,c) (e2,d) = IsZero $ normalizeExpr $ c *^ e2 ^-^ d *^ e1

normalizeExpr :: ExprZ -> ExprZ
normalizeExpr e = LA.mapCoeff (`div` d) e
  where d = abs $ gcd' $ map fst $ LA.terms e

-- "subst1Constr x t c" computes c[t/x]
subst1Constr :: Var -> LA.Expr Rational -> Constr -> Constr
subst1Constr x t c =
  case c of
    IsPos e    -> IsPos (f e)
    IsNonneg e -> IsNonneg (f e)
    IsZero e   -> IsZero (f e)
  where
    f :: ExprZ -> ExprZ
    f = normalizeExpr . fst . toRat . LA.applySubst1 x t . LA.mapCoeff fromInteger

-- | Simplify conjunction of 'Constr's.
-- It returns 'Nothing' when a inconsistency is detected.
simplify :: [Constr] -> Maybe [Constr]
simplify = fmap concat . mapM f
  where
    f :: Constr -> Maybe [Constr]
    f c@(IsPos e) =
      case LA.asConst e of
        Just x -> guard (x > 0) >> return []
        Nothing -> return [c]
    f c@(IsNonneg e) =
      case LA.asConst e of
        Just x -> guard (x >= 0) >> return []
        Nothing -> return [c]
    f c@(IsZero e) =
      case LA.asConst e of
        Just x -> guard (x == 0) >> return []
        Nothing -> return [c]

instance Eval (Model Rational) Constr Bool where
  eval m (IsPos t)    = LA.eval m (LA.mapCoeff fromInteger t :: LA.Expr Rational) > 0
  eval m (IsNonneg t) = LA.eval m (LA.mapCoeff fromInteger t :: LA.Expr Rational) >= 0
  eval m (IsZero t)   = LA.eval m (LA.mapCoeff fromInteger t :: LA.Expr Rational) == 0

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

fromLAAtom :: LA.Atom Rational -> DNF Constr
fromLAAtom (OrdRel a op b) = atomR' op (toRat a) (toRat b)
  where
    atomR' :: RelOp -> Rat -> Rat -> DNF Constr
    atomR' op a b = 
      case op of
        Le -> DNF [[a `leR` b]]
        Lt -> DNF [[a `ltR` b]]
        Ge -> DNF [[a `geR` b]]
        Gt -> DNF [[a `gtR` b]]
        Eql -> DNF [[a `eqR` b]]
        NEq -> DNF [[a `ltR` b], [a `gtR` b]]

toLAAtom :: Constr -> LA.Atom Rational
toLAAtom (IsNonneg e) = LA.mapCoeff fromInteger e .>=. LA.constant 0
toLAAtom (IsPos e)    = LA.mapCoeff fromInteger e .>.  LA.constant 0
toLAAtom (IsZero e)   = LA.mapCoeff fromInteger e .==. LA.constant 0

constraintsToDNF :: [LA.Atom Rational] -> DNF Constr
constraintsToDNF = andB . map fromLAAtom

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

{-
(ls1,ls2,us1,us2) represents
{ x | ∀(M,c)∈ls1. M/c≤x, ∀(M,c)∈ls2. M/c<x, ∀(M,c)∈us1. x≤M/c, ∀(M,c)∈us2. x<M/c }
-}
type Bounds = ([Rat], [Rat], [Rat], [Rat])

evalBounds :: Model Rational -> Bounds -> Interval Rational
evalBounds model (ls1,ls2,us1,us2) =
  Interval.intersections $
    [ Finite (evalRat model x) <=..< PosInf | x <- ls1 ] ++
    [ Finite (evalRat model x) <..<  PosInf | x <- ls2 ] ++
    [ NegInf <..<= Finite (evalRat model x) | x <- us1 ] ++
    [ NegInf <..<  Finite (evalRat model x) | x <- us2 ]

boundsToConstrs :: Bounds -> Maybe [Constr]
boundsToConstrs  (ls1, ls2, us1, us2) = simplify $ 
  [ x `leR` y | x <- ls1, y <- us1 ] ++
  [ x `ltR` y | x <- ls1, y <- us2 ] ++ 
  [ x `ltR` y | x <- ls2, y <- us1 ] ++
  [ x `ltR` y | x <- ls2, y <- us2 ]

collectBounds :: Var -> [Constr] -> (Bounds, [Constr])
collectBounds v = foldr phi (([],[],[],[]),[])
  where
    phi :: Constr -> (Bounds, [Constr]) -> (Bounds, [Constr])
    phi constr@(IsNonneg t) x = f False constr t x
    phi constr@(IsPos t) x = f True constr t x
    phi constr@(IsZero t) (bnd@(ls1,ls2,us1,us2), xs) =
      case LA.extractMaybe v t of
        Nothing -> (bnd, constr : xs)
        Just (c,t') -> ((t'' : ls1, ls2, t'' : us1, us2), xs)
          where
            t'' = (signum c *^ negateV t', abs c)

    f :: Bool -> Constr -> ExprZ -> (Bounds, [Constr]) -> (Bounds, [Constr])
    f strict constr t (bnd@(ls1,ls2,us1,us2), xs) =
      case LA.extract v t of
        (c,t') ->
          case c `compare` 0 of
            EQ -> (bnd, constr : xs)
            GT ->
              if strict
              then ((ls1, (negateV t', c) : ls2, us1, us2), xs) -- 0 < cx + M ⇔ -M/c <  x
              else (((negateV t', c) : ls1, ls2, us1, us2), xs) -- 0 ≤ cx + M ⇔ -M/c ≤ x
            LT ->
              if strict
              then ((ls1, ls2, us1, (t', negate c) : us2), xs) -- 0 < cx + M ⇔ x < M/-c
              else ((ls1, ls2, (t', negate c) : us1, us2), xs) -- 0 ≤ cx + M ⇔ x ≤ M/-c

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

{-| @'project' x φ@ returns @[(ψ_1, lift_1), …, (ψ_m, lift_m)]@ such that:

* @⊢_LRA ∀y1, …, yn. (∃x. φ) ↔ (ψ_1 ∨ … ∨ φ_m)@ where @{y1, …, yn} = FV(φ) \\ {x}@, and

* if @M ⊧_LRA ψ_i@ then @lift_i M ⊧_LRA φ@.
-}
project :: Var -> [LA.Atom Rational] -> [([LA.Atom Rational], Model Rational -> Model Rational)]
project v xs = do
  ys <- unDNF $ constraintsToDNF xs
  (zs, mt) <- maybeToList $ project' v ys
  return (map toLAAtom zs, mt)

project' :: Var -> [Constr] -> Maybe ([Constr], Model Rational -> Model Rational)
project' v cs = projectN' (IS.singleton v) cs

{-| @'projectN' {x1,…,xm} φ@ returns @[(ψ_1, lift_1), …, (ψ_n, lift_n)]@ such that:

* @⊢_LRA ∀y1, …, yp. (∃x. φ) ↔ (ψ_1 ∨ … ∨ φ_n)@ where @{y1, …, yp} = FV(φ) \\ {x1,…,xm}@, and

* if @M ⊧_LRA ψ_i@ then @lift_i M ⊧_LRA φ@.
-}
projectN :: VarSet -> [LA.Atom Rational] -> [([LA.Atom Rational], Model Rational -> Model Rational)]
projectN vs xs = do
  ys <- unDNF $ constraintsToDNF xs
  (zs, mt) <- maybeToList $ projectN' vs ys
  return (map toLAAtom zs, mt)

projectN' :: VarSet -> [Constr] -> Maybe ([Constr], Model Rational -> Model Rational)
projectN' = f
  where
    f vs xs
      | IS.null vs = return (xs, id)
      | Just (v,vdef,ys) <- findEq vs xs = do
          let mt1 m = IM.insert v (evalRat m vdef) m
          (zs, mt2) <- f (IS.delete v vs) [subst1Constr v (fromRat vdef) c | c <- ys]
          return (zs, mt1 . mt2)
      | otherwise = 
          case IS.minView vs of
            Nothing -> return (xs, id) -- should not happen
            Just (v,vs') -> 
              case collectBounds v xs of
                (bnd, rest) -> do
                  cond <- boundsToConstrs bnd
                  let mt1 m =
                       case Interval.simplestRationalWithin (evalBounds m bnd) of
                         Nothing  -> error "ToySolver.FourierMotzkin.project': should not happen"
                         Just val -> IM.insert v val m
                  (ys, mt2) <- f vs' (rest ++ cond)
                  return (ys, mt1 . mt2)

findEq :: VarSet -> [Constr] -> Maybe (Var, Rat, [Constr])
findEq vs = msum . map f . pickup
  where
    pickup :: [a] -> [(a,[a])]
    pickup [] = []
    pickup (x:xs) = (x,xs) : [(y,x:ys) | (y,ys) <- pickup xs]

    f :: (Constr, [Constr]) -> Maybe (Var, Rat, [Constr])
    f (IsZero e, cs) = do
      let vs2 = IS.intersection vs (vars e)
      guard $ not $ IS.null vs2
      let v = IS.findMin vs2
          (c, e') = LA.extract v e
      return (v, (negateV e', c), cs)
    f _ = Nothing

-- | @'solve' {x1,…,xn} φ@ returns @Just M@ that @M ⊧_LRA φ@ when
-- such @M@ exists, returns @Nothing@ otherwise.
--
-- @FV(φ)@ must be a subset of @{x1,…,xn}@.
-- 
solve :: VarSet -> [LA.Atom Rational] -> Maybe (Model Rational)
solve vs cs = msum [solve' vs cs2 | cs2 <- unDNF (constraintsToDNF cs)]

solve' :: VarSet -> [Constr] -> Maybe (Model Rational)
solve' vs cs = do
  (cs2,mt) <- projectN' vs =<< simplify cs
  let m :: Model Rational
      m = IM.empty
  guard $ all (eval m) cs2
  return $ mt m

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

gcd' :: [Integer] -> Integer
gcd' [] = 1
gcd' xs = foldl1' gcd xs

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