module Exercise_9 where

{-WETT-}
import Data.List (find, sortBy, (\\))
import Data.Ord (comparing)
import Data.Maybe (fromJust)
import Control.Arrow (first, second)

import qualified Data.IntMap as Map

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

{-T9.3.2-}
top :: Literal
top = Pos T

bottom :: Literal
bottom = Neg T

{-T9.3.3-}
clauseToForm :: Clause -> Form
clauseToForm [] = L bottom
clauseToForm ls = foldr ((:|:) . L) (L $ last ls) (init ls)

conjToForm :: ConjForm -> Form
conjToForm [] = L top
conjToForm ds = foldr ((:&:) . clauseToForm) (clauseToForm $ last ds) (init ds)

{-T9.3.4-}
substLiteral :: Valuation -> Literal -> Literal
substLiteral v l@(Pos (Var n)) = case lookup n v of
  Just b -> if b then top else bottom
  Nothing -> l
substLiteral v l@(Neg (Var n)) = case lookup n v of
  Just b -> if b then bottom else top
  Nothing -> l
substLiteral _ l = l

substClause :: Valuation -> Clause -> Clause
substClause = map . substLiteral

substConj :: Valuation -> ConjForm -> ConjForm
substConj = map . substClause

{-H9.2.1-}
simpConj :: ConjForm -> ConjForm
simpConj cs
  | any null cs' = [[]]
  | otherwise = cs'
  where
    cs' = simplify cs
    simplify = map remBottoms . filter noTop
    noTop = (top `notElem`)
    remBottoms = filter (/= bottom)

{-H9.2.2-}
cnf :: Form -> ConjForm
cnf (L literal) = [[literal]]
cnf (formA :&: formB) = cnf formA ++ cnf formB
cnf (formA :|: formB) = concatMap (\a -> map (\b -> a ++ b) cnfB) cnfA
  where
    cnfA = cnf formA
    cnfB = cnf formB

{-H9.2.3-}
selectV :: ConjForm -> Maybe (Name, Bool)
selectV clauses = assignVar <$> firstLiteralWithVar
  where
    firstLiteralWithVar = find literalIsVar $ concat clausesByLength
    clausesByLength = sortBy (comparing length) clausesWithVars
    clausesWithVars = map (filter literalIsVar) clauses
    literalIsVar (Pos a) = atomIsVar a
    literalIsVar (Neg a) = atomIsVar a
    assignVar (Pos a) = (atomToVarName a, True)
    assignVar (Neg a) = (atomToVarName a, False)
    atomIsVar T = False
    atomIsVar (Var _) = True
    atomToVarName (Var name) = name

{-H9.2.5-}
type IntLiteral = Int
type IntConjForm = [[IntLiteral]]

type IntAssignment = Int
type IntValuation = [IntAssignment]

type VarFrequencies = [IntLiteral]

type NPos = Int
type NNeg = Int
type VarMap = Map.IntMap (NPos, NNeg)

satConj :: ConjForm -> Maybe Valuation
satConj cs = fromIntValuation . addMissingAssignments <$> solveDpll
  where
    solveDpll = satConjDpll [] varFs varMap intCs'
    -- Unfortunately, the tests on the submission system require a valuation to
    -- also contain assignments of variables occurring exclusively in
    -- tautological clauses (even though their value does not matter), whereas
    -- omitting assignments of variables from clauses containing (Pos T) at the
    -- beginning seems to be fine...
    addMissingAssignments val = val ++ (Map.keys varMap \\ map abs val)
    -- Variable names (as ints), ordered by frequency, descending.
    varFs = map fst
      . sortBy (flip $ comparing snd)
      -- Calculate total occurrences.
      . map (second $ uncurry (+))
      . Map.toList $ varMap
    varMap = generateVarMap intCs
    intCs' = removeTautologicalClauses intCs
    intCs = toIntConjForm cs'
    cs' = simpConj' cs

simpConj' :: ConjForm -> ConjForm
simpConj' cs
  | any null cs' = [[]]
  | otherwise = cs'
  where
    cs' = map (filter (/= Neg T)) $ filter (Pos T `notElem`) cs

-- Convert literals to ints, e.g. (Neg (Var "11")) -> -11.
toIntConjForm :: ConjForm -> IntConjForm
toIntConjForm = map (map toInt)
  where
    toInt (Pos (Var name)) = read name
    toInt (Neg (Var name)) = -(read name)

generateVarMap :: IntConjForm -> VarMap
generateVarMap cs = foldl incrVM Map.empty (concat cs)
  where
    incrVM vM l = Map.alter incr (abs l) vM
      where
        incr Nothing = Just $ modifier (0, 0)
        incr (Just cur) = Just $ modifier cur
        modifier = if l > 0 then first succ else second succ

-- Remove clauses containing var as well as ¬var.
removeTautologicalClauses :: IntConjForm -> IntConjForm
removeTautologicalClauses = filter (not . anyTautologies)
  where
    anyTautologies ls = any (\l -> -l `elem` ls) ls

satConjDpll :: IntValuation -> VarFrequencies -> VarMap -> IntConjForm -> Maybe IntValuation
satConjDpll val varFs varM cs
  | null cs = Just val
  | any null cs = Nothing
  | otherwise = case assignUnitClauseVariable cs of
    -- Unit clause found -> unit propagation.
    Just u -> satConjDpll
      (u : val)
      (filter (/= abs u) varFs)
      varM
      (simpConjEvaluating u cs)
    -- No unit clauses -> backtracking.
    Nothing -> case satConjDpll (a : val) varFs' varM cs' of
      Nothing -> satConjDpll (-a : val) varFs' varM negCs'
      Just val' -> Just val'
      where
        (a, varFs') = assign varFs varM
        cs' = simpConjEvaluating a cs
        negCs' = simpConjEvaluating (-a) cs

simpConjEvaluating :: IntAssignment -> IntConjForm -> IntConjForm
simpConjEvaluating i = map (filter (/= negate i)) . filter (i `notElem`)

{-# INLINE simpConjEvaluating #-}

assignUnitClauseVariable :: IntConjForm -> Maybe IntAssignment
assignUnitClauseVariable cs
  | null ucs = Nothing
  | otherwise = Just $ head $ head ucs
  where
    ucs = filter lengthIsOne cs
    lengthIsOne [] = False
    lengthIsOne [_] = True
    lengthIsOne _ = False

{-# INLINE assignUnitClauseVariable #-}

assign :: VarFrequencies -> VarMap -> (IntAssignment, [Int])
assign (f : fs) m = (if nPos > nNeg then key else -key, fs)
  where
    (nPos, nNeg) = fromJust $ Map.lookup key m
    key = abs f

{-# INLINE assign #-}

fromIntValuation :: IntValuation -> Valuation
fromIntValuation = map fromInt
  where
    fromInt i = (show (abs i), i > 0)

{-H9.2.6-}
sat :: Form -> Maybe Valuation
sat = satConj . cnf
{-TTEW-}