module SAT.CDCL (cdcl, SATResult(..)) where

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

data SATResult
  = SAT (Map.Map Int Bool)
  | UNSAT [Clause]
  deriving (Show)

data DecisionLevel = DecisionLevel
  { dlVar :: Int
  , dlVal :: Bool
  , dlLevel :: Int
  , dlImplied :: [(Int, Bool)]
  } deriving (Show)

cdcl :: CNF -> SATResult
cdcl (CNF clauses numVars) = go Map.empty clauses [] 0
  where
    go asgn cls trail level =
      case unitPropagate asgn cls of
        Conflict ->
          if level == 0 then UNSAT cls
          else
            let learned = analyzeConflict trail cls
                newLevel = backtrackLevel trail learned
                asgn' = backtrack asgn trail newLevel
                trail' = take newLevel trail
            in go asgn' (learned : cls) trail' newLevel
        Propagated asgn' cls' ->
          if null cls' then SAT asgn'
          else
            let var = chooseVar asgn' cls'
                level' = level + 1
                asgn'' = Map.insert var True asgn'
                dl = DecisionLevel var True level' []
            in go asgn'' cls' (trail ++ [dl]) level'

analyzeConflict :: [DecisionLevel] -> [Clause] -> Clause
analyzeConflict trail _ =
  case trail of
    [] -> []
    dls -> let lastDl = last dls
           in [Neg (dlVar lastDl)]

backtrackLevel :: [DecisionLevel] -> Clause -> Int
backtrackLevel trail _ = max 0 (length trail - 1)

backtrack :: Map.Map Int Bool -> [DecisionLevel] -> Int -> Map.Map Int Bool
backtrack asgn trail level =
  let toRemove = drop level trail
      vars = map dlVar toRemove
  in foldr Map.delete asgn vars

chooseVar :: Map.Map Int Bool -> [Clause] -> Int
chooseVar asgn clauses =
  let allVars = nub $ map litVar $ concat clauses
      unassigned = filter (\v -> not (Map.member v asgn)) allVars
  in case unassigned of
       (v:_) -> v
       [] -> 0