module Exercise09 where

import Data.Bifunctor as B
import Data.List as L
import Data.Map as M
import Data.Maybe
import Data.Ord (Down (Down))
import Data.Set as S
import Types

-- 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 (Lit _ n) = n

lPos :: Name -> Literal
lPos = Lit Pos

lNeg :: Name -> Literal
lNeg = Lit Neg

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

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

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

lENegate :: ELiteral -> ELiteral
lENegate (Literal Pos n) = Literal Neg n
lENegate (Literal Neg n) = Literal Pos n

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

evalLiteral :: Valuation -> Literal -> Bool
evalLiteral val l@(Lit _ n)
  | lIsPos l = n `elem` val
  | otherwise = n `notElem` val

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

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

clauseIsTauto :: Clause -> Bool
clauseIsTauto = clauseIsTautoE . toEClause

clauseIsTautoE :: EClause -> Bool
clauseIsTautoE [] = False
clauseIsTautoE (l : ls) = lENegate l `elem` ls || clauseIsTautoE ls

-- homework starts here

-- feel free to change behaviour and type
resolve :: Name -> Clause -> Clause -> Clause
resolve n cp cn = S.delete (Lit Pos n) cp `S.union` S.delete (Lit Neg n) cn

-- feel free to change behaviour and type

resolvants :: KeyClause -> Map Literal (Set KeyClause) -> (([Resolve], [Clause]), Map Literal (Set KeyClause))
resolvants (Kc k uc) m = S.foldr (\l ((rs, cs), n) -> (S.foldr (\(Kc j c) (ss, ds) -> if lIsPos l then (Res (lName l) k j : ss, resolve (lName l) uc c : ds) else (Res (lName l) j k : ss, resolve (lName l) c uc : ds)) (rs, cs) (findWithDefault S.empty (lNegate l) m), insertWith S.union l (S.singleton (Kc k uc)) n)) (([], []), m) uc

-- proof checking

-- don't change the type nor expected behaviour
proofCheck :: EConjForm -> Proof -> Bool
proofCheck cf (Ref ((Res n k l) : rs)) = Lit Pos n `elem` pos && Lit Neg n `elem` neg && proofCheck (cf ++ [toEClause (resolve n pos neg)]) (Ref rs)
  where
    pos = toClause (cf !! k)
    neg = toClause (cf !! l)
proofCheck cf (Ref _) = [] `elem` cf
proofCheck cf (Mod val) = eval val (toConjForm cf)

-- feel free to change behaviour and type
selClause :: SelClauseStrategy
selClause = S.deleteFindMin

-- feel free to change behaviour and type
resolutionParam :: Int -> Set KeyClause -> Map Literal (Set KeyClause) -> (Proof, KeyConjForm)
resolutionParam n cs m
  | S.null cs = (Mod [], [])
  | S.null (getClause (fst c)) = (Ref [], S.elems cs)
  | clauseIsTauto (getClause (fst c)) || fst c `elem` findWithDefault S.empty (minimum ((getClause . fst) c)) m = resolutionParam n (uncurry filterSubsumed c) m
  | otherwise = (if fst p /= Mod [] then Ref ((fst . fst) res ++ fromMaybe [] (getRes (fst p))) else Mod [], fst c : snd p)
  where
    c = selClause cs
    p = resolutionParam (n + length ((fst . fst) res)) (filterSubsumed (fst c) (toKeySet n ((snd . fst) res) `S.union` snd c)) (snd res)
    res = resolvants (fst c) m

filterSubsumed :: KeyClause -> Set KeyClause -> Set KeyClause
filterSubsumed clause = S.filter (not . isSubsetOf (getClause clause) . getClause)

getRes :: Proof -> Maybe [Resolve]
getRes (Ref rs) = Just rs
getRes (Mod _) = Nothing

toKeySet :: Int -> ConjForm -> Set KeyClause
toKeySet n = S.fromList . toKeyList n

toKeyList :: Int -> ConjForm -> [KeyClause]
toKeyList n = L.zipWith Kc (iterate (+ 1) n)

toConjForm2 :: Set KeyClause -> ConjForm
toConjForm2 = L.map (\(Kc _ c) -> c) . S.elems

{-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 cs = (if fst res == Mod [] then Mod (extractModel (nub (L.map getClause (snd res)))) else fst res, snd res)
  where
    res = resolutionParam len set M.empty
    len = length cs
    set = (toKeySet 0 . toConjForm) cs

{-TTEW-}

-- 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 :: ConjForm -> Valuation
extractModel = run [] . sort . Prelude.map (sortOn Down) . toEConjForm
  where
    run val [] = val
    run val (c : cs)
      | evalClause val (toClause c) = run val cs
      | otherwise = run ((lName . toLiteral) (head c) : val) cs