module Exercise09 where

import Types
import FTrie

import Data.Function (on)
import qualified Data.IntSet as IS
import qualified Data.IntMap.Lazy as IM
import Control.Arrow (second)
import Data.List (sortBy, sortOn)

-- util functions

lPos :: Name -> Literal
lPos = id

lNeg :: Name -> Literal
lNeg = negate

lIsPos :: Literal -> Bool
lIsPos = (>0)

lIsNeg :: Literal -> Bool
lIsNeg = (<0)

lNegate :: Literal -> Literal
lNegate = negate

lName :: Literal -> Name
lName = abs

-- evaluations

evalLiteral :: Valuation -> Literal -> Bool
evalLiteral val l
  | lIsPos l = l `IS.member` val
  | otherwise = lNegate l `IS.notMember` val

evalClause :: Valuation -> Clause -> Bool
evalClause val = IS.foldr ((||) . evalLiteral val) False

eval :: Valuation -> ConjForm -> Bool
eval = all . evalClause

clauseIsTauto :: Clause -> Bool
clauseIsTauto c = IS.foldr ((||) . (`IS.member` c) . lNegate) False c

-- basic resolution functions

complements :: Literal -> Literal -> Bool
complements n n' = n == -n'

resolve :: Name -> Clause -> Clause -> Clause
resolve n cp cn = IS.delete (lPos n) cp `IS.union` IS.delete (lNeg n) cn

-- proof checking

proofCheck :: EConjForm -> Proof -> Bool
proofCheck cs (Model' val) = eval val (toConjForm cs)
proofCheck initECs (Refutation' initRs) = run (IM.size initCs) initCs initRs
  where 
    initCs = toConjForm initECs
    run _ cs [] = IS.empty `elem` cs
    run i cs (Resolve' n k1 k2:rs) = run (i+1) (IM.insert i c cs) rs
      where c = resolve n (cs IM.! k1) (cs IM.! k2)

-- homework starts here

-- selection strategies

selClause :: SelClauseStrategy 
selClause q
  | IM.null q = Nothing 
  | otherwise = let (k, kdcs) = IM.findMin q -- pick the smallest clause
                    (kdc, kdcs') = IM.deleteFindMin kdcs in
    Just (kdc, IM.update (const $ queueMaybeChild kdcs') k q)

queueMaybeChild :: KeyDataConjForm -> Maybe KeyDataConjForm
queueMaybeChild kdcs
  | IM.null kdcs = Nothing
  | otherwise  = Just kdcs

selLiterals :: SelLiteralsStrategy
selLiterals ((mi,_),_)
  | lIsNeg mi = IS.singleton mi
  | otherwise = IS.empty

resolvants :: SelLiteralsStrategy -> KeyDataClause -> KeyDataClause -> [(Resolve, Clause)]
resolvants selLiterals kdc1@(_,dc1) kdc2@(_,dc2) =
  let sel1 = selLiterals dc1
      sel2 = selLiterals dc2 in
  produce kdc1 sel1 kdc2 sel2 ++ produce kdc2 sel2 kdc1 sel1
  where
    produce (kp,((_,lmp),cp)) selp (kn,((_,lmn),cn)) seln
      | lIsNeg lmp || not (IS.null selp) = []
      | IS.null seln = [rc | lmp `complements` lmn]
      | otherwise = [rc | lNegate lmp `IS.member` seln]
      where rc = let n = lName lmp in (Resolve' n kp kn, resolve n cp cn)

isTrivial :: Clause -> Bool
isTrivial = clauseIsTauto

preProcess :: ConjForm -> KeyConjForm
preProcess = IM.filter (not . isTrivial) -- remove trivial clauses

resolutionParam :: SelClauseStrategy -> SelLiteralsStrategy -> ConjForm -> (Proof, KeyConjForm)
resolutionParam selClause selLiterals csInit = run [] (IM.size csInit) qInit IM.empty emptyFTrie
  where
    preProcessed = preProcess csInit
    qInit = keyConjFormToQueue preProcessed
    run :: [Resolve] -> Key -> Queue -> KeyDataConjForm -> FTrie -> (Proof, KeyConjForm)
    run rs k uqkdcs pkdcs pFTrie = case selClause uqkdcs of -- pick next clause
      -- everything is saturated -> extract the model
      Nothing -> let pkcs = IM.map snd pkdcs in (Model' $ extractModel pkcs, pkcs)
      Just (ukdc@(uk,udc@(_,uc)), uqkdcs') -> 
        -- found the empty clause -> return the resolution proof
        if IS.null uc then (Refutation' rs, IM.map snd pkdcs `IM.union` queueToKeyConjForm uqkdcs')
        else if pFTrie `fTrieSubsumes` uc then run rs k uqkdcs' pkdcs pFTrie -- clause subsumed -> continue
        else 
          let (pkdcs', pFTrie') = filterSubsumedDataKeyConjFormFTrie uc pkdcs pFTrie -- remove clauses subsumed by uc
              (nrs, ncs) = unzip
                 [(r, c) | (r, c) <- concatMap (resolvants selLiterals ukdc) $ IM.toList pkdcs', -- get all candidates
                           not $ isTrivial c] -- skip trivial clauses
              nuqkdcs = zip [k..] ncs `addToQueue` uqkdcs' -- add new clauses to queue
              pkdcs'' = IM.insert uk udc pkdcs' -- add clause to processed set
              pFTrie'' = insertFTrie uk uc pFTrie' in -- add clause to FTrie
          run (rs ++ nrs) (k + length ncs) nuqkdcs pkdcs'' pFTrie'' -- continue
    addToQueue :: [(Key, Clause)] -> Queue -> Queue
    addToQueue kcs q = foldr (\(k,c) -> IM.insertWith IM.union (IS.size c) $ IM.singleton k (lData c, c)) q kcs

resolution :: EConjForm -> (Proof, KeyConjForm)
resolution = resolutionParam selClause selLiterals . toConjForm

lCompare :: Literal -> Literal -> Ordering
lCompare l1 l2
  | abs l1 == abs l2 = compare l2 l1
  | otherwise = (compare `on` abs) l1 l2

-- as described on https://lara.epfl.ch/w/_media/sav08/gbtalk.pdf
extractModel :: ConjForm -> Valuation
extractModel cs = run IS.empty $ sortOn snd $ map (second $ sortBy (flip lCompare) . IS.toList) $ IM.toList cs
  where
    run :: Valuation -> [(Key, [Literal])] -> Valuation
    run val [] = val
    run val ((k,c):kcs)
      | evalClause val (cs IM.! k) = run val kcs
      | otherwise = run (IS.insert (head c) val) kcs