module Exercise09 where

import Data.HashSet as HS
import Data.List (nub, sort, sortOn)
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 (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 `elem` val
  | otherwise = n `notElem` val

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

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

clauseIsTauto :: Clause -> Bool
clauseIsTauto c = any (\x -> HS.member (lNegate x) c) c

-- homework starts here

-- feel free to change behaviour and type
resolve :: Name -> Clause -> Clause -> Clause
resolve n c1 c2 = HS.union c1New c2New
  where
    c1New = HS.delete (Literal Pos n) c1
    c2New = HS.delete (Literal Neg n) c2

-- feel free to change behaviour and type
resolvants :: KeyClause -> KeyClause -> [(Resolve, Clause)]
resolvants kc1 kc2 = resolvantsHelper kc1 kc2 ++ resolvantsHelper kc2 kc1

resolvantsHelper :: KeyClause -> KeyClause -> [(Resolve, Clause)]
resolvantsHelper (id1, c1) (id2, c2) = [(Resolve n id1 id2, resolve n c1 c2) | n <- HS.toList literals] -- todo: ineffiecient
  where
    positive = HS.filter lIsPos c1
    inverted = HS.map lNegate positive
    exists = HS.filter (`elem` inverted) c2
    literals = HS.map lName exists

-- don't change the type nor expected behaviour
proofCheck :: EConjForm -> Proof -> Bool
proofCheck cnf (Refutation res) = proofCheckHelper (toKeyMap 0 (toConjForm cnf)) res
proofCheck cnf (Model valuation) = eval valuation (toConjForm cnf)

proofCheckHelper :: KeyConjForm -> [Resolve] -> Bool
proofCheckHelper kcnf res
  | isNothing allClauses = False
  | otherwise = HS.fromList [] `elem` S.toList (fromJust allClauses)
  where
    allClauses = proofCheckHelperAllClauses kcnf res

proofCheckHelperAllClauses :: KeyConjForm -> [Resolve] -> Maybe ConjForm
proofCheckHelperAllClauses kcnf [] = Just (S.fromDistinctAscList (M.elems kcnf))
proofCheckHelperAllClauses kcnf (r : rs)
  | isNothing nextKcnf = Nothing
  | otherwise = proofCheckHelperAllClauses (fromJust nextKcnf) rs
  where
    nextKcnf = insertNewClause kcnf r

insertNewClause :: KeyConjForm -> Resolve -> Maybe KeyConjForm
insertNewClause kcnf (Resolve n k1 k2)
  | isNothing val1 || isNothing val2 = Nothing
  | otherwise = Just (insertNewClauseHelper kcnf (M.size kcnf, resolve n (fromJust val1) (fromJust val2)))
  where
    val1 = M.lookup k1 kcnf
    val2 = M.lookup k2 kcnf

insertNewClauseHelper :: KeyConjForm -> KeyClause -> KeyConjForm
insertNewClauseHelper kcnf clause = uncurry M.insert clause kcnf

-- feel free to change behaviour and type
selClause :: SelClauseStrategy
selClause kcnf = findNextSmallest kcnf 0

---- Always selects the shortest clause
findNextSmallest :: KeyConjForm -> Int -> Maybe KeyClause
findNextSmallest kcnf n
  | M.empty == kcnf = Nothing
  | M.empty == fst partitioned = findNextSmallest kcnf (n + 1)
  | otherwise = Just (head (M.toList (fst partitioned)))
  where
    partitioned = M.partition (\x -> length x == n) kcnf

-- feel free to change behaviour and type
resolutionParam :: SelClauseStrategy -> ConjForm -> (Proof, KeyConjForm)
resolutionParam str cnf
  | S.member (HS.fromList []) cnf = (Refutation [], toKeyMap 0 cnf)
  | otherwise = resolutionParamsStepDoNext str filteredMap M.empty hashset offset []
  where
    filteredMap = smallerAndEqualKcnf cnf
    offset = S.size cnf - M.size filteredMap --- in order to create valid proofs we must remember how many clauses we discarded
    hashset = HS.fromList (S.toList cnf) --- we want to remember all the clauses we discarded, so later we dont add them again

{--
Creates a smaller and equal formula by:
- removing tautologies
- removing clauses that are subsumed by others
- removing duplicates
--}
smallerAndEqualKcnf :: ConjForm -> KeyConjForm
smallerAndEqualKcnf cnf = filteredClauses keyMapND keyMapND
  where
    keyMap = toKeyMap 0 cnf
    keyMapND = nubKeyConjForm keyMap

isSubsumed :: Clause -> KeyConjForm -> Bool
isSubsumed c1 kcnf = any (\x -> isSubset x c1 && x /= c1) elements
  where
    elements = HS.fromList (M.elems kcnf)

-- true if c1 is subset of c2
isSubset :: Clause -> Clause -> Bool
isSubset c1 c2 = all (`HS.member` c2) c1

{--
str: SelClauseStrategy
c : The clasue we are currently looking at
u : unprocessed clauses
p: processed clauses
original: the initial clauses (needed so we don't look at the same clauses again)
offset: the number of discarded claueses (needed to create a correct proof)
r : resolution steps up to now
--}
resolutionParamsStep :: SelClauseStrategy -> KeyClause -> KeyConjForm -> KeyConjForm -> HashSet Clause -> Int -> [Resolve] -> (Proof, KeyConjForm)
resolutionParamsStep str c u p original offset r
  | Prelude.null (snd c) = (Refutation r, M.union pn un)
  | otherwise = resolutionParamsStepDoNext str un pn (HS.union original (HS.fromList (snd rs))) (offset + updatedOffset) rn
  where
    rs = getResolvants c p
    pn = uncurry M.insert c p
    newClauses = S.fromDistinctAscList (snd rs)
    newClausesMap = toKeyMap (M.size p + M.size u + offset + 1) newClauses
    newClausesNotInOriginal = removeContained (nubKeyConjForm newClausesMap) original
    newClausesFiltered = filteredClauses (u `M.union` pn) newClausesNotInOriginal
    updatedOffset = M.size newClausesMap - M.size newClausesFiltered
    un = u `M.union` newClausesFiltered
    rn = r ++ fst rs

removeContained :: KeyConjForm -> HashSet Clause -> KeyConjForm
removeContained kcnf original = M.filter (\x -> not (HS.member x original)) kcnf

filteredClauses :: KeyConjForm -> KeyConjForm -> KeyConjForm
filteredClauses kcnf newClauses = filtered
  where
    filtered = M.filter (\x -> not (clauseIsTauto x || isSubsumed x kcnf)) newClauses

-- removes duplicates, leaving the clasue with the highest key
nubKeyConjForm :: KeyConjForm -> KeyConjForm
nubKeyConjForm kcnf = M.fromAscList (removeDuplicates pairs)
  where
    pairs = M.assocs kcnf

removeDuplicates :: [(Int, Clause)] -> [(Int, Clause)]
removeDuplicates [] = []
removeDuplicates (x : xs)
  | any (\(_, r) -> r == snd x) xs = rest
  | otherwise = x : rest
  where
    rest = removeDuplicates xs

toKeyMap :: Int -> ConjForm -> KeyConjForm
toKeyMap n cnf = M.fromAscList (zip [n ..] (S.toList cnf))

resolutionParamsStepDoNext :: SelClauseStrategy -> KeyConjForm -> KeyConjForm -> HashSet Clause -> Int -> [Resolve] -> (Proof, KeyConjForm)
resolutionParamsStepDoNext str u p original offset r = resolutionParamsStepDoNextCheckMaybe str nextClause u p original offset r
  where
    nextClause = str u

resolutionParamsStepDoNextCheckMaybe :: SelClauseStrategy -> Maybe KeyClause -> KeyConjForm -> KeyConjForm -> HashSet Clause -> Int -> [Resolve] -> (Proof, KeyConjForm)
resolutionParamsStepDoNextCheckMaybe _ Nothing _ p _ _ _ = (Model (extractModel (S.fromDistinctAscList (M.elems p))), p)
resolutionParamsStepDoNextCheckMaybe str (Just c) u p original offset r = resolutionParamsStep str c (M.delete (fst c) u) p original offset r

getResolvants :: KeyClause -> KeyConjForm -> ([Resolve], [Clause])
getResolvants c kcnf = unzip (concat allResolvants)
  where
    allResolvants = M.elems (M.mapWithKey (curry (resolvants c)) kcnf)

{-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 cnf = resolutionParam selClause (toConjForm cnf)

{-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 cnf = run [] (toExtractable cnf)
  where
    run val [] = val
    run val (c : cs)
      | evalEClause val c = run val cs
      | otherwise = run (lName (head c) : val) cs

evalEClause :: Valuation -> [ELiteral] -> Bool
evalEClause = any . evalLiteral

toExtractable :: ConjForm -> [[ELiteral]]
toExtractable cnf = sort (Prelude.map (sortOn Down) list)
  where
    list = toEConjForm cnf