-- module to parse CNF (Dimacs) format of problems and solutions

module Cnf
  ( parseProblem
  , parseSolution
  , myParse
  , CNFProblem(..)
  , Solution(..)
  , Trail
  , Clause
  , Clauses
  ) where

import Control.Applicative ((<*), (*>))
import Text.ParserCombinators.Parsec hiding (spaces)
import Data.Set (Set, fromList)

-- data structures

data Solution =
  Sat | Unsat | Unknown
  deriving (Show, Eq)

data CNFProblem = CNFProblem
  { numClauses :: Int
  , numVars    :: Int
  , clauses    :: [Clause]
  } deriving (Show, Eq)

type Trail = Set Int
type Clause = Set Int
type Clauses = [Clause]

-- parse function
myParse :: Parser a -> String -> a
myParse parser input = either (error . show) id $ parse parser "" input

-- parse problem

parseProblem :: Parser CNFProblem
parseProblem = do
  problem <- skipMany commentLine *> problemLine
  clauses <- skipMany commentLine *> (clause `endBy` nspaces) <* eof
  return $ problem { clauses = map fromList clauses }

commentLine :: Parser Char
commentLine = char 'c' >> skipMany (noneOf "\n") >> newline

problemLine :: Parser CNFProblem
problemLine = do
  numVars <- (char 'p' >> spaces >> string "cnf" >> spaces) *> int
  numClauses <- spaces *> int <* nspaces
  return $ CNFProblem numClauses numVars []

clause :: Parser [Int]
clause = do
  literal <- int
  if literal == 0
    then return []
    else do -- recurse (read the rest of clause)
      c <- nspaces *> clause
      return $ literal : c

int :: Parser Int
int = do
  neg <- optionMaybe (char '-')
  digits <- many1 digit
  return $ (if neg == Nothing then 1 else (-1)) * read digits

spaces = skipMany (oneOf " \t")
nspaces = skipMany (oneOf " \t\n")

-- parse solution
parseSolution :: Parser (Solution, Trail)
parseSolution = do
  skipMany commentLine >> char 's' >> spaces
  sol <- (string "SATISFIABLE" <|> string "UNSATISFIABLE" <|> string "UNKNOWN")
  skipMany (oneOf " \t\n")
  trail <- readTrail <|> return []
  return (readSolution sol, fromList trail)

readTrail :: Parser [Int]
readTrail = do
  i <- skipMany (oneOf "v \t\n") *> int
  if i == 0
    then return []
    else do
      trail <- readTrail
      return $ i : trail

-- read a solution
readSolution :: String -> Solution
readSolution "SATISFIABLE"    = Sat
readSolution "UNSATISFIABLE"  = Unsat
readSolution "UNKNOWN"        = Unknown