module Exercise09 where

import Types

import Data.List (sort, sortOn)
import Data.Ord (Down(Down))
import qualified Data.HashSet as Set
import qualified Data.Sequence as Seq
import Test.QuickCheck

-- things from the tutorial

instance Show Polarity where
    show Pos = ""
    show Neg = "~"

instance Show ELiteral where
    show (Literal p a) = show p ++ a

lName :: Literal -> Name
lName (Literal _ n) = n

lPos :: Name -> Literal
lPos = Literal Pos

lNeg :: Name -> Literal
lNeg = Literal Neg

lIsPos :: Literal -> Bool
lIsPos (Literal p _) = p == Pos

lIsNeg :: Literal -> Bool
lIsNeg = not . lIsPos

lNegate :: Literal -> Literal
lNegate (Literal Pos n) = Literal Neg n
lNegate (Literal Neg n) = Literal Pos n

complements :: Literal -> Literal -> Bool
complements (Literal p n) (Literal p' n') = p /= p' && n == n'

evalLiteral :: Valuation -> Literal -> Bool
evalLiteral val l@(Literal _ n)
    | lIsPos l  = n `Set.member` val
    | otherwise = not $ n `Set.member` val

evalClause :: Valuation -> Clause -> Bool
evalClause = any . evalLiteral

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

clauseIsTauto :: Clause -> Bool
clauseIsTauto = clauseIsTauto' . toEClause
    where
        clauseIsTauto' []       = False
        clauseIsTauto' (l:ls)   = lNegate l `elem` ls || clauseIsTauto' ls




-- homework starts here

-- feel free to change behaviour and type
resolve :: Name -> Clause -> Clause -> Clause
resolve n cp cn = (lPos n `Set.delete` cp) `Set.union` (lNeg n `Set.delete` cn)

-- feel free to change behaviour and type
resolvants :: KeyClause -> KeyClause -> Set.HashSet (Resolve, Clause)
resolvants clause1 clause2 = resolvants' clause1 clause2 `Set.union`
                             resolvants' clause2 clause1
    where
        resolvants' (KeyClause k1 c1) (KeyClause k2 c2) =
            Set.map (\literal -> (Resolve (lName literal) k1 k2,
                                  resolve (lName literal) c1 c2))
                    literals
            where
                literals = Set.filter (\lit -> lIsPos lit &&
                                               lNegate lit `Set.member` c2)
                                      c1

-- proof checking

-- don't change the type nor expected behaviour
proofCheck :: EConjForm -> Proof -> Bool
proofCheck clauses (Mod val)    = eval val (toConjForm clauses)
proofCheck clauses (Ref steps)  = check (Seq.fromList (map toClause clauses)) steps ([] `elem` clauses)
    where
        check _       []                                unsatisfiable = unsatisfiable
        check clauses (Resolve name key1 key2 : steps)  unsatisfiable =
            let
                clause1 = clauses `Seq.index` key1
                clause2 = clauses `Seq.index` key2
                newClause = resolve name clause1 clause2
            in
                validStep name clause1 clause2 &&
                check (clauses Seq.|> newClause) steps (unsatisfiable || null newClause)
        
        validStep name clause1 clause2 =
            lPos name `Set.member` clause1 &&
            lNeg name `Set.member` clause2


-- feel free to change behaviour and type
selClause :: SelClauseStrategy 
selClause keyClauses
    | null keyClauses   = Nothing
    | otherwise         = Just $ head $ sortOn (\(KeyClause _ clause) -> Set.size clause) $ Set.toList keyClauses  -- TODO better strategy

-- feel free to change behaviour and type
resolutionParam :: SelClauseStrategy -> EConjForm -> (Proof, KeyConjForm)
resolutionParam strategy clauses = loop (Set.fromList (zipWith (\clause key -> KeyClause key (toClause clause)) clauses [0..]))
                                        Set.empty [] (length clauses)
    where
        loop :: Set.HashSet KeyClause -> Set.HashSet KeyClause -> [Resolve] -> Int -> (Proof, KeyConjForm)
        loop unprocessed processed steps nextKey =
            case strategy unprocessed of
                Nothing ->
                    (Mod $ toValuation (extractModel $ toEConjForm (Set.map (\(KeyClause _ clause) -> clause) processed)), processed)
                Just keyClause@(KeyClause _ clause) ->
                    if null clause then
                        (Ref steps, unprocessed `Set.union` processed)
                    else
                        let
                            resolvants              = getAllResolvants keyClause processed
                            newResolvants           = filter (\(_, clause) -> not $ KeyClause 0 clause `Set.member` processed) resolvants
                            (newSteps, newClauses)  = (map fst newResolvants,
                                                       Set.fromList (zipWith KeyClause [nextKey..] (map snd newResolvants)))
                        in
                            loop ((keyClause `Set.delete` unprocessed) `Set.union` newClauses)
                                 (keyClause `Set.insert` processed)
                                 (steps ++ newSteps)
                                 (nextKey + Set.size newClauses)

        -- resolvants :: KeyClause -> KeyClause -> Set.HashSet (Resolve, Clause)
        getAllResolvants :: KeyClause -> Set.HashSet KeyClause -> [(Resolve, Clause)]
        getAllResolvants keyClause otherKeyClauses =
            Set.toList $
            Set.unions (map (resolvants keyClause) (Set.toList otherKeyClauses))


{-WETT-}

-- don't change the type; instantiate your best resolution prover here
-- Please leave some comments for the MC Jr to understand your ideas and code :')
resolution :: EConjForm -> (Proof, KeyConjForm)
resolution = resolutionParam selClause

{-TTEW-}

prop_resolution eConjForm = within 1000000 $ proofCheck eConjForm (fst (resolution eConjForm))
-- prop_resolution eConjForm = proofCheck eConjForm (fst (resolution eConjForm))

-- extract a model as described on https://lara.epfl.ch/w/_media/sav08/gbtalk.pdf

instance Ord Polarity where
    Neg <= Pos = False
    _ <= _ = True

instance Ord ELiteral where
    (Literal p a) <= (Literal p' a') = a < a' || a == a' && p <= p'


extractModel :: EConjForm -> EValuation
extractModel = run [] . sort . map (sortOn Down)
    where
        run val [] = val
        run val (c:cs)
            | evalClause (toValuation val) (toClause c) = run val cs
            | otherwise = run (lName (head c):val) cs