module Exercise09 where

import Types

import Data.List
import Data.Maybe
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.Set (Set)
import qualified Data.Set as Set

-- 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 = abs

lPos :: Name -> Literal
lPos = abs

-- TODO: could be optimized
lNeg :: Name -> Literal
lNeg = negate . lPos

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

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

lNegate :: Literal -> Literal
lNegate = negate

complements :: Literal -> Literal -> Bool
complements l l' = lName l == lName l' && l /= l'

evalClause :: Valuation -> Clause -> Bool
evalClause val clause
  | IntSet.null clause = False
  | otherwise          = check pos_literals || check neg_literals
    where check = not . IntSet.null
          pos_literals = IntSet.intersection val clause
          -- Check if there are negated literals which are not set to true by
          -- this valuation. (i.e., they would be true in the end :-) )
          neg_literals = (IntSet.map lName . IntSet.filter lIsNeg) clause IntSet.\\ val

eval :: Valuation -> ConjForm -> Bool
eval val = foldr ((&&) . evalClause val) True

clauseIsTauto :: Clause -> Bool
clauseIsTauto clause = (not . IntSet.null) $ IntSet.intersection clause (IntSet.map lNegate clause)

-- homework starts here

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

-- feel free to change behaviour and type
resolvants :: KeyClause -> KeyClause -> [(Resolve, Clause)]
resolvants (KeyClause k1 c1) (KeyClause k2 c2) = resolvants''
    -- inline lIsPos and lName
    where (pos1, neg1) = IntSet.partition (> 0) c1
          (pos2, neg2) = IntSet.partition (> 0) c2
          -- all literals that are pos in c1 and neg in c2
          to_resolve1 = IntSet.intersection pos1 $ IntSet.map abs neg2
          -- all literals that are neg in c1 and pos in c2
          to_resolve2 = IntSet.intersection pos2 $ IntSet.map abs neg1
          -- Note the order of the clauses which are passed to "resolve"!
          -- Also note the order of the keys passed to "Resolve"!
          resolvants' = IntSet.foldr (\n list ->
                            (MyResolve n k1 k2, resolve n c1 c2) : list) [] to_resolve1
          resolvants'' = IntSet.foldr (\n list ->
                            (MyResolve n k2 k1, resolve n c2 c1) : list) resolvants' to_resolve2

-- proof checking

-- don't change the type nor expected behaviour
proofCheck :: EConjForm -> Proof -> Bool
-- Explicitly handle the trivial case.
proofCheck [] _ = True
-- Search for the empty clause if the resolution list is empty.
proofCheck ecf (MyRefutation []) = [] `elem` ecf
proofCheck ecf (MyRefutation res_list) = proofCheckInner (length ecf) (toConjForm ecf) res_list
    where proofCheckInner _ _ [] = False
          proofCheckInner new_index cf ((MyResolve n k k'):rs)
            | k >= new_index || k' > new_index = False -- indices are greater than current map size
            | IntSet.null new_clause = True
            | otherwise = proofCheckInner (new_index + 1) (IntMap.insert new_index new_clause cf) rs
            where c = fromJust $ cf IntMap.!? k
                  c' = fromJust $ cf IntMap.!? k'
                  new_clause = resolve n c c'
proofCheck ecf (MyModel val) = eval val $ toConjForm ecf

-- feel free to change behaviour and type

-- Note: As this is an internal function, we do not need the Maybe stuff here
-- because it's guaranteed that the KeyConjForm is not empty.
-- Return the KeyClause with the smallest size (= the smallest key of our KeyConjForm).
selClause :: SelClauseStrategy 
selClause = (\(SizeKeyClause _ key_clause, new_kcf) -> (key_clause, new_kcf)) . Set.deleteFindMin

-- feel free to change behaviour and type
resolutionParam :: SelClauseStrategy -> KeyConjForm -> KeyConjForm -> [Resolve] -> Int -> (Proof, KeyConjForm)
-- "num_clauses" represents the next index (= key) to be used for new keyclauses.
-- Note that we have to keep the clause numbering!
resolutionParam selKeyClause unprocessed processed res_steps num_clauses
  | Set.null unprocessed = (MyModel (extractModel processed), Set.union unprocessed processed)
  | IntSet.null selected_clause = (MyRefutation (reverse res_steps), Set.union unprocessed processed)
--  | (\(KeyClause key _) -> key) key_clause > 100 = (MyRefutation res_steps, Set.union unprocessed processed)
  | otherwise = resolutionParam selKeyClause unprocessed' processed' res_steps' num_clauses'
--  | otherwise = resolutionParam selKeyClause unprocessed' processed' res_steps' num_clauses'
  where
      -- Select smallest clause available.
      (key_clause, unprocessed'') = selKeyClause unprocessed
      selected_clause = (\(KeyClause _ clause) -> clause) key_clause
      -- Get a List of (Resolve, Clause).
      res_list = Set.foldr (\(SizeKeyClause _ kc) list -> resolvants key_clause kc ++ list) [] processed
      -- Expand the set of resolution steps.
      -- Enumerate the new resolution steps in order to preserve the correct
      -- resolution order.
      -- Add the new resolution steps in the wrong order. (We do not want ++)
      res_steps' = foldl (\list (r, _) -> r : list) res_steps res_list
      -- Get new total clause number.
      num_clauses' = num_clauses + length res_list
      -- Add the new clauses to the new unprocessed Map.
      -- Use "num_clauses" as index to start the enumerating of the new clauses.
      -- Filter tautologies **after** enumerating the new clauses! Otherwise,
      -- the proof may get uncompatible.
      new_unprocessed = Set.filter (\(SizeKeyClause _ (KeyClause _ c)) -> not $ clauseIsTauto c)
                        $ Set.fromList
                        $ zipWith (\(_, clause) key ->
                            SizeKeyClause (IntSet.size clause) (KeyClause key clause))
                            res_list [num_clauses..num_clauses']
      -- Add our selected clause to the processed clauses.
      new_clause = SizeKeyClause (IntSet.size selected_clause) key_clause
      processed'' = if new_clause `isNotSuperSet` processed then
                        Set.insert new_clause processed
                   else processed
      -- Actually merge the new clauses.
      unprocessed' = Set.union unprocessed''
                     $ Set.filter (\c -> c `isNotSuperSet` unprocessed''
                                      && c `isNotSuperSet` processed'')
                       new_unprocessed Set.\\ processed''
      processed' = keyConjFormRemoveSuperSets unprocessed' processed''

{-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 ecf = resolutionParam selClause kcf Set.empty [] (length ecf)
    -- Convert the EConjForm to a KeyConjForm. (Preserve the original indices!)
    where kcf = keyConjFormRemoveOwnSuperSets
                $ keyConjFormRemoveTautos
                $ conjForm2KeyConjForm
                $ toConjForm ecf

-- TODO: is this lazy?
isNotSuperSet :: SizeKeyClause -> KeyConjForm -> Bool
isNotSuperSet skc@(SizeKeyClause _ (KeyClause _ clause)) =
    Set.null . Set.filter (\c' -> skc /= c' && clause `isSuperSetOf` c')
        where isSuperSetOf c (SizeKeyClause _ (KeyClause _ c')) = c' `IntSet.isSubsetOf` c

keyConjFormRemoveTautos :: KeyConjForm -> KeyConjForm
keyConjFormRemoveTautos = Set.filter (\(SizeKeyClause _ (KeyClause _ c)) -> not $ clauseIsTauto c)

-- Remove the clauses in the second kcf that are supersets of clauses in the
-- first kcf.
keyConjFormRemoveSuperSets :: KeyConjForm -> KeyConjForm -> KeyConjForm
keyConjFormRemoveSuperSets kcf = Set.filter (`isNotSuperSet` kcf)

keyConjFormRemoveOwnSuperSets :: KeyConjForm -> KeyConjForm
keyConjFormRemoveOwnSuperSets kcf = keyConjFormRemoveSuperSets kcf kcf

{-TTEW-}

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

-- For every clause add one of the non-negated variables
-- Sort the clauses in descending order, so that the negated literals come as
-- late as possible.
extractModel :: KeyConjForm -> Valuation
extractModel = run IntSet.empty
               . map (\(SizeKeyClause _ (KeyClause _ c)) -> IntSet.toDescList c)
               . Set.toAscList
    where run val [] = val
          run val (c:cs)
            | evalClause val $ IntSet.fromList c = run val cs
            | otherwise = run (IntSet.insert (lName $ head c) val) cs