{-# LANGUAGE FlexibleInstances
           , TypeSynonymInstances #-}

module System.Expression ( Type(..)
                         , Expression(..)
                         , Relation(..)
                         , Formula(..)
                         , VariableName
                         , VariableVersion
                         , int
                         , constValue
                         , varType
                         , varName
                         , varVersion
                         , var
                         , var'
                         , var0
                         , varAt
                         , intVar
                         , intVar'
                         , intVar0
                         , intVarAt
                         , vars
                         , types
                         , substitute
                         , Predicate
                         , extractPredicates
                         , and
                         , or
                         , not
                         , atoms
                         , literals ) where

import Prelude hiding (and, or, not)

import Data.Map (Map)
import Data.Monoid

import qualified Data.List as L
import qualified Data.List.Ordered as O
import qualified Data.Map as M

import Utils.Canonical
import Utils.Versioned

type VariableName    = String
type VariableVersion = Int

class WithSubstitutableExpression v where
    vars       :: v -> [Expression]
    substitute :: Map Expression Expression -> v -> v

class Typed v where
    types :: v -> [Type]

data Type = Integer | Container Type deriving (Eq, Ord, Show, Read)

int :: Type
int = Integer

data Expression = Var !Type !VariableName !VariableVersion
                | Const !Int
                | Add [Expression]
                | Mul [Expression]
    deriving (Eq, Ord, Show, Read)

constValue :: Expression -> Int
constValue (Const c) = c

varType :: Expression -> Type
varType (Var t _ _) = t

varName :: Expression -> VariableName
varName (Var _ v _) = v

varVersion :: Expression -> VariableVersion
varVersion (Var _ _ n) = n

var :: Type -> VariableName -> Expression
var t n = Var t n 0

var' :: Type -> VariableName -> Expression
var' t n = Var t n 1

var0 :: Type -> VariableName -> Expression
var0 t = var t . (++ "0")

varAt :: Type -> VariableVersion -> VariableName -> Expression
varAt t = flip (Var t)

intVar :: VariableName -> Expression
intVar = var int

intVar' :: VariableName -> Expression
intVar' = var' int

intVar0 :: VariableName -> Expression
intVar0 = var0 int

intVarAt :: VariableVersion -> VariableName -> Expression
intVarAt = varAt int

instance WithSubstitutableExpression Expression where
    vars v@Var {}   = [v]
    vars   (Add as) = L.nub $ concatMap vars as
    vars   (Mul ms) = L.nub $ concatMap vars ms
    vars _          = []

    substitute m e = M.findWithDefault d e m where
        d = case e of
            v@Var {}    -> v
            c@(Const _) -> c
            (Add as)    -> Add $ map (substitute m) as
            (Mul ms)    -> Mul $ map (substitute m) ms

instance Typed Expression where
    types (Const c)   = [Integer]
    types (Var t _ _) = [t]
    types (Add as)    = concatMap types as
    types (Mul ms)    = concatMap types ms

instance Normal Expression where
    normalise c@(Const _) = c
    normalise v@Var {} = v
    normalise (Add as) = inline1
                       . Add
                       . L.sort
                       . constants 0
                       . inline
                       . map normalise

                       $ as where
        inline :: [Expression] -> [Expression]
        inline as = inline' as []

        inline' :: [Expression] -> [Expression] -> [Expression]
        inline' []             out = out
        inline' (Add as' : as) out = inline' as $ inline' as' out
        inline' (a       : as) out = inline' as $ a : out

        inline1 :: Expression -> Expression
        inline1   (Add [])  = Const 0
        inline1   (Add [a]) = a
        inline1 a@(Add _)   = a
        inline1   _         = error "unexpected expression"

        constants :: Int -> [Expression] -> [Expression]
        constants c [] | c == 0     = []
                       | otherwise  = [Const c]
        constants c (Const c' : es) = constants (c + c') es
        constants c (e        : es) = e : constants c es
    normalise (Mul ms) = inline1
                       . Mul
                       . L.sort
                       . constants 1
                       . inline
                       . map normalise

                       $ ms where
        inline :: [Expression] -> [Expression]
        inline ms = inline' ms []

        inline' :: [Expression] -> [Expression] -> [Expression]
        inline' []             out = out
        inline' (Mul ms' : ms) out = inline' ms $ inline' ms' out
        inline' (m       : ms) out = inline' ms $ m : out

        inline1 :: Expression -> Expression
        inline1   (Mul [])  = Const 1
        inline1   (Mul [m]) = m
        inline1 m@(Mul _)   = m
        inline1   _         = error "unexpected expression"

        constants :: Int -> [Expression] -> [Expression]
        constants c [] | c == 1     = []
                       | otherwise  = [Const c]
        constants c (Const c' : es) = constants (c * c') es
        constants c (e        : es) = e : constants c es

instance Canonical Expression where
    canonicalise (Var t v _) = var t v
    canonicalise e = normalise e

instance Versioned Expression where
    bumped (Const _)   = fromVersions 0 []
    bumped (Var _ v n) = fromVersions 0 [(v, n)]
    bumped (Add as)    = mergeVersionsWith max 0 . map bumped $ as
    bumped (Mul ms)    = mergeVersionsWith max 0 . map bumped $ ms

    bumped1 (Const _)   = fromVersions 0 []
    bumped1 (Var _ v _) = fromVersions 0 [(v, 1)]
    bumped1 (Add as)    = mergeVersionsWith max 0 . map bumped $ as
    bumped1 (Mul ms)    = mergeVersionsWith max 0 . map bumped $ ms

    bump _ c@(Const _) = c
    bump m (Var t v n) = Var t v $ n + getVersion m v
    bump m (Add as)    = Add $ map (bump m) as
    bump m (Mul ms)    = Mul $ map (bump m) ms

    invert _ c@(Const _) = c
    invert m (Var t v n) = Var t v $ max (getVersion m v) n - n
    invert m (Add as)    = Add $ map (invert m) as
    invert m (Mul ms)    = Mul $ map (invert m) ms

    reset c@(Const _)   = c
    reset   (Var t v _) = var t v
    reset   (Add as)    = Add $ map reset as
    reset   (Mul ms)    = Mul $ map reset ms

data Relation = Eq | Ne | Gt | Ge | Lt | Le
    deriving (Eq, Ord, Show, Read)

data Formula = Atom Expression !Relation Expression
             | Not Formula
             | And [Formula]
             | Or  [Formula]
    deriving (Eq, Show, Read)

atoms :: Formula -> [Formula]
atoms f = normalise $ atoms' f [] where
    atoms' a@Atom {}   as = a : as
    atoms'   (Not f)   as = atoms' f as
    atoms'   (And fs)  as = foldr atoms' as fs
    atoms'   (Or  fs)  as = foldr atoms' as fs

literals :: Formula -> [Formula]
literals f = normalise $ literals' f [] where
    literals' a@Atom {}       as = a : as
    literals' n@(Not Atom {}) as = n : as
    literals'   (Not f)       as = literals' f as
    literals'   (And fs)      as = foldr literals' as fs
    literals'   (Or  fs)      as = foldr literals' as fs

instance WithSubstitutableExpression Formula where
    vars (Atom a _ b) = L.nub $ vars a ++ vars b
    vars (Not f)      = vars f
    vars (And as)     = L.nub $ concatMap vars as
    vars (Or os)      = L.nub $ concatMap vars os

    substitute m (Atom a p b) = Atom (substitute m a) p (substitute m b)
    substitute m (Not f)      = Not $ substitute m f
    substitute m (And as)     = And $ map (substitute m) as
    substitute m (Or os)      = Or  $ map (substitute m) os

instance Typed Formula where
    types (Atom a _ b) = types a ++ types b
    types (Not f)      = types f
    types (And as)     = concatMap types as
    types (Or os)      = concatMap types os

instance Ord Formula where
    (Not f1) `compare` (Not f2) = f1 `compare` f2 <> EQ
    f1       `compare` (Not f2) = f1 `compare` f2 <> LT
    (Not f1) `compare` f2       = f1 `compare` f2 <> GT

    (Atom a b c) `compare` (Atom d e f) =
        a `compare` d <> b `compare` e <> c `compare` f

    Atom {} `compare` _       = LT
    _       `compare` Atom {} = GT

    (And a1) `compare` (And a2) = a1 `compare` a2
    (And _)  `compare` _        = LT
    _        `compare` (And _)  = GT

    (Or o1) `compare` (Or o2) = o1 `compare` o2

type Predicate = Formula

-- can be put into equivalent normal form
instance Normal Formula where
    normalise (Atom a Eq b) =
        let (a', b') = normalise (a, b) in
            Atom a' Eq b'
    normalise (Atom a Ne b) = normalise (Atom a Eq b)
    normalise (Atom a Gt b) = simplify $ Atom (normalise b) Lt (normalise a)
    normalise (Atom a Ge b) = Not
                            . simplify $ Atom (normalise a) Lt (normalise b)
    normalise (Atom a Lt b) = simplify $ Atom (normalise a) Lt (normalise b)
    normalise (Atom a Le b) = Not
                            . simplify $ Atom (normalise b) Lt (normalise a)
    normalise (Not f) =
        case normalise f of
            Not f' -> f'
            And [] -> Or  []
            Or  [] -> And []
            f'     -> propagateNot . Not $ f'
    normalise (And as) = inlineAnd1
                       . And
                       . O.nub
                       . L.sort
                       . mergeEq
                       . inlineAnd
                       . map normalise

                       $ as
    normalise (Or os) = inlineOr1
                      . Or
                      . O.nub
                      . L.sort
                      . inlineOr
                      . map normalise

                      $ os

-- can be put into canonical form, phase can change
instance Canonical Formula where
    canonicalise (Atom a Eq b) =
        let (a', b') = canonicalise (a, b) in
            Atom a' Eq b'
    canonicalise (Atom a Ne b) = canonicalise (Atom a Eq b)
    canonicalise (Atom a Gt b) = simplify $ Atom (canonicalise b) Lt (canonicalise a)
    canonicalise (Atom a Ge b) = simplify $ Atom (canonicalise a) Lt (canonicalise b)
    canonicalise (Atom a Lt b) = simplify $ Atom (canonicalise a) Lt (canonicalise b)
    canonicalise (Atom a Le b) = simplify $ Atom (canonicalise b) Lt (canonicalise a)
    canonicalise (Not f) = canonicalise f
    canonicalise (And [a]) = canonicalise a
    canonicalise (And as) = inlineAnd1
                          . And
                          . O.nub
                          . L.sort
                          . inlineAnd
                          . map canonicalise

                          $ as
    canonicalise (Or [o]) = canonicalise o
    canonicalise (Or os) = inlineOr1
                         . Or
                         . O.nub
                         . L.sort
                         . inlineOr
                         . map canonicalise

                         $ os

instance Versioned Formula where
    bumped (Atom a p b) = mergeVersionsWith max 0 [bumped a, bumped b]
    bumped (Not f)      = bumped f
    bumped (And as)     = mergeVersionsWith max 0 . map bumped $ as
    bumped (Or os)      = mergeVersionsWith max 0 . map bumped $ os

    bumped1 (Atom a p b) = mergeVersionsWith max 0 [bumped1 a, bumped1 b]
    bumped1 (Not f)      = bumped1 f
    bumped1 (And as)     = mergeVersionsWith max 0 . map bumped1 $ as
    bumped1 (Or os)      = mergeVersionsWith max 0 . map bumped1 $ os

    bump m (Atom a p b) = Atom (bump m a) p (bump m b)
    bump m (Not f)      = Not (bump m f)
    bump m (And as)     = And . map (bump m) $ as
    bump m (Or os)      = Or  . map (bump m) $ os

    invert m (Atom a p b) = Atom (invert m a) p (invert m b)
    invert m (Not f)      = Not (invert m f)
    invert m (And as)     = And . map (invert m) $ as
    invert m (Or os)      = Or  . map (invert m) $ os

    reset (Atom a p b) = Atom (reset a) p (reset b)
    reset (Not f)      = Not (reset f)
    reset (And as)     = And . map reset $ as
    reset (Or os)      = Or  . map reset $ os

inlineAnd :: [Formula] -> [Formula]
inlineAnd as = inline' as [] where

    inline' :: [Formula] -> [Formula] -> [Formula]
    inline' []             out = out
    inline' (Or []   : _ ) _   = [Or []]
    inline' (And as' : as) out = inline' as $ inline' as' out
    inline' (a       : as) out = inline' as $ a : out

inlineAnd1 :: Formula -> Formula
inlineAnd1   (And [a]) = a
inlineAnd1 a@(And _)   = a
inlineAnd1   _         = error "unexpected formula"

mergeEq :: [Formula] -> [Formula]
mergeEq as = m as as where
    m [] _ = []
    m (a@(Not (Atom x Lt y)) : as) bs =
        let r = m as bs
            e = normalise (Atom x Eq y) in if Not (Atom y Lt x) `elem` bs then e : r else a : r
    m (a : as) bs = a : m as bs

inlineOr :: [Formula] -> [Formula]
inlineOr os = inline' os [] where

    inline' :: [Formula] -> [Formula] -> [Formula]
    inline' []            out = out
    inline' (And [] : _ ) _   = [And []]
    inline' (Or os' : os) out = inline' os $ inline' os' out
    inline' (o      : os) out = inline' os $ o : out

inlineOr1 :: Formula -> Formula
inlineOr1   (Or [o]) = o
inlineOr1 o@(Or _)   = o
inlineOr1   _        = error "unexpected formula"

simplify :: Formula -> Formula
simplify (Atom (Add as1) p (Add as2)) =
        let (ns1, ps1) = L.partition isNeg as1
            (ns2, ps2) = L.partition isNeg as2 in
            Atom (normalise . Add $ ps1 ++ map neg ns2) p (normalise . Add $ ps2 ++ map neg ns1)
simplify (Atom a p (Add as)) = let (ns, ps) = L.partition isNeg as in
    Atom (normalise . Add $ a : map neg ns) p (normalise $ Add ps)
simplify (Atom (Add as) p b) = let (ns, ps) = L.partition isNeg as in
    Atom (normalise $ Add ps) p (normalise . Add $ b : map neg ns)
simplify (Atom a p b) | isNeg a = Atom (Const 0) p (normalise $ Add [neg a, b])
simplify (Atom a p b) | isNeg b = Atom (normalise $ Add [a, neg b]) p (Const 0)
simplify f = f

neg :: Expression -> Expression
neg e = Mul [Const (-1), e]

isNeg :: Expression -> Bool
isNeg (Mul [Const (-1), _]) = True
isNeg (Mul [_, Const (-1)]) = True
isNeg _                     = False

and :: Formula -> Formula -> Formula
and a b = normalise $ And [a, b]

or :: Formula -> Formula -> Formula
or a b = normalise $ Or [a, b]

not :: Formula -> Formula
not a = normalise $ Not a

extractPredicates :: Formula -> [Predicate]
extractPredicates f = map reset $ extractPredicates' f [] where
    extractPredicates' :: Formula -> [Predicate] -> [Predicate]
    extractPredicates' a@Atom {} ps = canonicalise a : ps
    extractPredicates' (Not f) ps   = extractPredicates' f ps
    extractPredicates' (And as) ps  = foldr extractPredicates' ps as
    extractPredicates' (Or os) ps   = foldr extractPredicates' ps os

propagateNot :: Formula -> Formula
propagateNot a@Atom {}        = a
propagateNot   (Not (Not f))  = propagateNot f
propagateNot n@(Not Atom {})  = n
propagateNot   (Not (And as)) = Or  $ map (propagateNot . Not) as
propagateNot   (Not (Or os))  = And $ map (propagateNot . Not) os
propagateNot   (And as)       = And $ map propagateNot as
propagateNot   (Or os)        = Or  $ map propagateNot os
