module Exercise_9 where 
{-WETT-}

import Data.List (foldl', find)
import Data.Maybe (fromJust, fromMaybe, isJust)
import Control.Applicative ((<|>))
import System.Environment (getArgs)
import Text.Read (readMaybe)
import qualified Data.IntMap.Strict as IM
import qualified Data.IntSet as IS

type Name = String
data Atom = T | Var Name
  deriving (Eq, Show, Read)
data Literal = Pos Atom | Neg Atom deriving (Eq, Show, Read)
data Form = L Literal | Form :&: Form | Form :|: Form
  deriving (Eq, Show)
type Clause = [Literal]
type ConjForm = [Clause]
type Valuation = [(Name, Bool)]

-- let's do a better encoding:
type Name' = Int
-- i encodes "Pos (Var "vi")" and -i encodes "Neg (Var "vi")"
-- top will be encodes as (maxBound :: Int) and bottom as -(maxBound :: Int)
type Literal' = Int
type Clause' = IS.IntSet -- use sets --> no redundancy!
type ConjForm' = [Clause']
type Valuation' = IM.IntMap Bool

substLiteral :: Valuation' -> Name' -> Literal' -> Literal'
substLiteral v top l = maybe l (\b -> (if b then id else negate) top * signum l) $ v IM.!? abs l

substClause :: Valuation' -> Literal' -> Clause' -> Clause'
substClause v = IS.map . substLiteral v

substConj :: Valuation' -> Literal' -> ConjForm' -> ConjForm'
substConj v = map . substClause v

{-H9.2.1-}
simpClause :: Name' -> Clause' -> Clause'
simpClause top c 
  | IS.member top c = IS.singleton top
  | otherwise = IS.delete (-top) c

simpConj :: Name' -> ConjForm' -> ConjForm'
simpConj top = foldl' collect []
  where
    collect acc _ | acc == singleEmptyConj = singleEmptyConj
    collect acc c = let c' = simpClause top c in
                    if (IS.size c') == 0 then singleEmptyConj
                    else if c' == singleTrueClause then acc
                    else c':acc
    singleEmptyConj  = [IS.empty]
    singleTrueClause = IS.singleton top

choice :: [Maybe a] -> a
choice = fromJust . fromJust . find isJust

assignLiteral :: Literal' -> (Name', Bool)
assignLiteral l = (abs l, signum l > 0)

-- primitive selection strategy: selects the largest variable in the first clause 
-- sets the found variable to True if it occurs positively and False otherwise
-- selectFirstMax :: ConjForm' -> Maybe (Name', Bool)
-- selectFirstMax (c:cs) = Just $ assignLiteral $ IS.findMax c

-- selects from the shortest clause the variable that occurs the most frequently in the formula
selectMostFrequentShort :: ConjForm' -> Maybe (Name', Bool)
selectMostFrequentShort (c:cs) = Just $ assignLiteral resL
  where
    ((sc,_), cm) = foldl' (\(s, cm) c -> (updateShortest s c, updateCountMap cm c)) ((c, IS.size c), updateCountMap IM.empty c) cs
    updateShortest s@(_,ss) c = if IS.size c < ss then (c, IS.size c) else s
    updateCountMap = IS.foldl' (\cm l -> IM.alter (maybe (Just 1) (Just . (+1))) (abs l) cm)
    count l  = cm IM.! abs l
    (resL,_) = IS.foldl' (\acc@(_,c) l' -> if count l' > c then (l', count l') else acc) (0,-1) sc

-- select next variable
select :: ConjForm' -> (Name', Bool)
-- select cs = choice [selectFirstMax cs]
select cs = choice [selectMostFrequentShort cs]

-- select all unit clauses
selectUnit :: ConjForm' -> Valuation'
selectUnit = foldl' insertUnits IM.empty 
  where
    insertUnits val c = if IS.size c == 1 then let m = (IS.findMax c) in IM.insert (abs m) (m>0) val
                        else val

-- select all clauses occuring with single polarity
selectPure :: ConjForm' -> Valuation'
selectPure = IM.foldlWithKey' (\val k v -> if v/=0 then IM.insert k (v>0) val else val) IM.empty . foldl' goClause IM.empty
  where
    goClause = IS.foldl' updatePures
    updatePures val l = IM.alter (\m -> (m >>= \p -> Just $ if signum p == signum l then p else 0) <|> Just l) (abs l) val

-- select all safe assignments
selectSafe :: ConjForm' -> Valuation'
selectSafe cs = selectUnit cs `IM.union` selectPure cs

vars :: ConjForm' -> IS.IntSet
vars = foldl' IS.union IS.empty

-- finishes if there is at least one positive or one negative literal in each clause
finishTrivial :: Valuation' -> ConjForm' -> Maybe Valuation'
finishTrivial v cs = case foldl' checkNegPos (True, True) cs of
  (_,True) -> Just (v `IM.union` IS.foldl' (\m l -> IM.insert (abs l) True m) IM.empty (vars cs))
  (True,_) -> Just (v `IM.union` IS.foldl' (\m l -> IM.insert (abs l) False m) IM.empty (vars cs))
  _        -> Nothing
  where
    checkNegPos (False, False) c = (False, False)
    checkNegPos (False, True)  c = (False, updatedPos c)
    checkNegPos (True, False)  c = (updatedNeg c, False)
    checkNegPos (True, True)   c = (updatedNeg c, updatedPos c)
    updatedNeg c = IS.findMin c < 0
    updatedPos c = IS.findMax c > 0

{-H9.2.5-}
satWithSelection :: (ConjForm' -> Valuation') -> (ConjForm' -> (Name', Bool)) -> Name' -> ConjForm' -> Maybe Valuation'
satWithSelection selSafe sel top cs = run IM.empty $ simpConj top cs
  where
    run v [] = Just v -- satisfied
    run v [c] | c == IS.empty = Nothing -- unsatisfiable
    run v f = finishTrivial v f -- finish if there is at least one positive or one negative literal in each clause
          <|> run newV newF -- set next variables as selected
          <|> if selected then run newV' newF' -- set negation of selection if needed
              else Nothing
      where
        sV = selSafe f -- select all variables whose valuation can safely be set
        selected = IM.size sV == 0
        (n,b) = sel f -- select assignment for next variable
        -- new valuation and formula with selected assignment of next variable
        newV  = (if selected then IM.insert n b v else sV `IM.union` v)
        newF  = simpConj top $ substConj newV top f
        -- new valuation and formula with the negation of selected assignment of next variable
        newV' = IM.insert n (not b) v
        newF' = simpConj top $ substConj newV' top f

satConj' :: Name' -> ConjForm' -> Maybe Valuation'
satConj' n cs = Nothing <|> satWithSelection selectSafe select n cs

toName' :: Name -> Name'
toName' = read

toLiteral' :: Name' -> Literal -> Literal'
toLiteral' top (Pos T) = top
toLiteral' top (Neg T) = -top
toLiteral' _ (Pos (Var n)) = toName' n
toLiteral' _ (Neg (Var n)) = -toName' n

toClause' :: Name' -> Clause -> Clause'
toClause' top = foldl' (\acc l-> IS.insert (toLiteral' top l) acc) IS.empty

toConj' :: Name' -> ConjForm -> ConjForm'
toConj' = map . toClause'

toName :: Name' -> Name
toName = show

toValuation :: Valuation' -> Valuation
toValuation = IM.foldlWithKey' (\acc k v -> (toName k,v):acc) []

satConj :: ConjForm -> Maybe Valuation
satConj cs = let top = maxBound in satConj' top (toConj' top cs) >>= (Just . toValuation)
-- satConj :: ConjForm -> Maybe Valuation'
-- satConj cs = let top = maxBound in satConj' top (toConj' top cs)
{-TTEW-}