module SAT.UnitProp (unitPropagate, UnitResult(..)) where

import SAT.CNF
import qualified Data.Map.Strict as Map
import Data.List (nub)

data UnitResult
  = Propagated (Map.Map Int Bool) [Clause]
  | Conflict
  deriving (Show)

unitPropagate :: Map.Map Int Bool -> [Clause] -> UnitResult
unitPropagate assignment clauses = go assignment clauses
  where
    go asgn cls =
      case findUnit cls asgn of
        Nothing ->
          let simplified = simplifyClauses asgn cls
          in if any null simplified
             then Conflict
             else Propagated asgn simplified
        Just (var, val) ->
          let asgn' = Map.insert var val asgn
              simplified = simplifyClauses asgn' cls
          in if any null simplified
             then Conflict
             else go asgn' simplified

findUnit :: [Clause] -> Map.Map Int Bool -> Maybe (Int, Bool)
findUnit clauses asgn = go clauses
  where
    go [] = Nothing
    go (c:cs) =
      let unresolved = filter (not . resolved asgn) c
      in case unresolved of
           [lit] -> Just (litVar lit, isPositive lit)
           _ -> go cs
    isPositive (Pos _) = True
    isPositive (Neg _) = False

resolved :: Map.Map Int Bool -> Literal -> Bool
resolved asgn (Pos v) = Map.member v asgn
resolved asgn (Neg v) = Map.member v asgn

simplifyClauses :: Map.Map Int Bool -> [Clause] -> [Clause]
simplifyClauses asgn = filter (not . satisfied asgn) . map (removeResolved asgn)

satisfied :: Map.Map Int Bool -> Clause -> Bool
satisfied asgn = any (satLit asgn)

satLit :: Map.Map Int Bool -> Literal -> Bool
satLit asgn (Pos v) = Map.lookup v asgn == Just True
satLit asgn (Neg v) = Map.lookup v asgn == Just False

removeResolved :: Map.Map Int Bool -> Clause -> Clause
removeResolved asgn = filter (\lit -> not (falseLit asgn lit))

falseLit :: Map.Map Int Bool -> Literal -> Bool
falseLit asgn (Pos v) = Map.lookup v asgn == Just False
falseLit asgn (Neg v) = Map.lookup v asgn == Just True