module Types where

import           Data.Bifunctor  (second)
import           Data.IntSet     (IntSet)
import           Data.Map.Strict (Map)
import           Data.Set        (Set)

import           Test.QuickCheck as QC

-- import qualified Data.Map.Strict as Map
import qualified Data.IntSet     as IntSet
-- import qualified Data.Set        as Set

-- type definitions
-- don't change the unmarked types

type EName = String
type Name = FName -- change as you like
data Polarity = Pos | Neg
  deriving Eq
data ELiteral = Literal Polarity EName
  deriving Eq
type Literal = FLiteral -- change as you like
type EClause = [ELiteral]
type Clause = FClause -- change as you like
type EConjForm = [EClause]
type ConjForm = [Clause] -- change as you like

type Key = Int
type EKeyClause = (Key, EClause)
type KeyClause = (Key, Clause) -- change as you like
type EKeyConjForm = [EKeyClause]
type KeyConjForm = [KeyClause] -- change as you like

data EResolve = Resolve EName Key Key
  deriving (Eq, Show)
data FResolve = FResolve Name Key Key -- Fast Resolve (Using faster types)
  deriving (Eq, Show)
type Resolve = FResolve -- change as you like
type EValuation = [EName]
type Valuation = [Name] -- change as you like
data EProof = Refutation [EResolve] | Model EValuation
  deriving (Eq, Show)
type Proof = FProof -- change as you like

type SelClauseStrategy = KeyConjForm -> Maybe KeyClause  -- change as you like

-- Aliases for more readable code
type UnitClauseKeys = Map Literal Key
type SolverStepData = (Key, KeyConjForm, KeyConjForm, FConjForm, [Resolve], Valuation)
type SolverStep = (SolverStepData -> Either Proof SolverStepData)
type WatchedLiterals = Map Key (Literal, Literal)

-- Faster Types

type FName = Int
type FLiteral = Int -- Encoding Pos & Neg as (+) and (-)
type FClause = IntSet -- Faster Clause using a Set for O(log n) member testing
type FConjForm = Set Clause -- Faster ConjForm using a Set
type FKeyConjForm = Map Key Clause -- Faster KeyConjForm using a Map
data FProof = FRefutation [Resolve] | FModel Valuation -- Fast Proof (Using faster types)
  deriving (Eq, Show)

-- Wrapper around the data the DPLL algorithm needs
data DPLLData = DPLLData { key :: Key, processed :: [Key], unprocessed :: [Key], valuation :: Valuation, rSteps :: [Resolve], clauses :: FConjForm, mapKeyClause :: FKeyConjForm, watched :: Map Key (Literal, Literal), unit :: FClause, unitKey :: Map Literal Key }
    deriving (Show)

isPos :: FLiteral -> Bool
isPos l = l > 0

isNeg :: FLiteral -> Bool
isNeg l = l < 0

-- adapt the following if you changed the internal data types above

toEName :: Name -> EName
toEName = show

toName :: EName -> Name
toName = read

toELiteral :: Literal -> ELiteral
toELiteral n
    | isPos n   = Literal Pos (show n)
    | otherwise = Literal Neg (show (-n))

toLiteral :: ELiteral -> Literal
toLiteral (Literal Pos l) = read l
toLiteral (Literal Neg l) = -(read l)

toEClause :: Clause -> EClause
toEClause = map toELiteral . IntSet.toList

toClause :: EClause -> Clause
toClause = IntSet.fromList . map toLiteral

toEConjForm :: ConjForm -> EConjForm
toEConjForm = map toEClause

toConjForm :: EConjForm -> ConjForm
toConjForm = map toClause

toEKeyClause :: KeyClause -> EKeyClause
toEKeyClause = second toEClause

toKeyClause :: EKeyClause -> KeyClause
toKeyClause = second toClause

toEKeyConjForm :: KeyConjForm -> EKeyConjForm
toEKeyConjForm = map toEKeyClause

toKeyConjForm :: EKeyConjForm -> KeyConjForm
toKeyConjForm = map toKeyClause

toEResolve :: Resolve -> EResolve
toEResolve (FResolve n k1 k2) = Resolve (toEName n) k1 k2

toResolve :: EResolve -> Resolve
toResolve (Resolve n k1 k2) = FResolve (toName n) k1 k2

toEValuation :: Valuation -> EValuation
toEValuation = map toEName

toValuation :: EValuation -> Valuation
toValuation = map toName

toEProof :: Proof -> EProof
toEProof (FRefutation rs) = Refutation (map toEResolve rs)
toEProof (FModel v)       = Model (toEValuation v)

toProof :: EProof -> Proof
toProof (Refutation rs) = FRefutation (map toResolve rs)
toProof (Model v)       = FModel (toValuation v)

-- Moved Orphan classes
instance Show Polarity where
  show Pos = ""
  show Neg = "~"

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

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

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

instance Arbitrary Polarity where
  arbitrary = oneof $ map return [Pos, Neg]

-- mV: maximum variable index
genEName :: Int -> Gen EName
genEName mV = do
  n <- choose (1, mV)
  return $ show n

instance Arbitrary ELiteral where
  arbitrary = sized $ \n -> do
    p <- arbitrary
    v <- genEName n
    return $ Literal p v
  shrink (Literal p n) = map (Literal p . show) ([1..(read n - 1)] :: [Int])