module Main where

import System.Exit (exitSuccess, exitFailure)
import qualified Data.Map.Strict as Map

import IR.Boolean
import IR.NAND
import IR.BitVec
import SAT.CNF
import SAT.UnitProp
import SAT.DPLL
import Proof.Certificate
import Proof.Produce
import Checker.Kernel
import Language.Parser
import Language.Elaborator

main :: IO ()
main = do
  results <- sequence
    [ test "NAND truth table" testNAND
    , test "NOT from NAND" testNOT
    , test "AND from NAND" testAND
    , test "OR from NAND" testOR
    , test "XOR from NAND" testXOR
    , test "Half adder sum" testHalfAdderSum
    , test "Half adder carry" testHalfAdderCarry
    , test "BitVec add 3+5=8" testBitVecAdd
    , test "Tseitin preserves satisfiability" testTseitin
    , test "DPLL finds SAT" testDPLLSat
    , test "DPLL finds UNSAT" testDPLLUnsat
    , test "Unit propagation" testUnitProp
    , test "Proof certificate valid" testProofValid
    , test "Checker accepts valid" testCheckerAccepts
    , test "Checker rejects invalid" testCheckerRejects
    , test "Parse module" testParse
    , test "Elaborate module" testElaborate
    , test "NAND normal form" testNANDNormal
    , test "De Morgan via eval" testDeMorgan
    , test "MUX correctness" testMux
    ]
  let passed = length (filter id results)
      total = length results
  putStrLn $ "\n" ++ show passed ++ "/" ++ show total ++ " tests passed."
  if passed == total then exitSuccess else exitFailure

test :: String -> Bool -> IO Bool
test name result = do
  putStrLn $ (if result then "[OK]  " else "[FAIL]") ++ " " ++ name
  return result

testNAND :: Bool
testNAND =
  eval Map.empty (BNand BTrue BTrue) == False &&
  eval Map.empty (BNand BTrue BFalse) == True &&
  eval Map.empty (BNand BFalse BTrue) == True &&
  eval Map.empty (BNand BFalse BFalse) == True

testNOT :: Bool
testNOT =
  eval Map.empty (bnot BTrue) == False &&
  eval Map.empty (bnot BFalse) == True

testAND :: Bool
testAND =
  eval Map.empty (band BTrue BTrue) == True &&
  eval Map.empty (band BTrue BFalse) == False &&
  eval Map.empty (band BFalse BTrue) == False &&
  eval Map.empty (band BFalse BFalse) == False

testOR :: Bool
testOR =
  eval Map.empty (bor BTrue BTrue) == True &&
  eval Map.empty (bor BTrue BFalse) == True &&
  eval Map.empty (bor BFalse BTrue) == True &&
  eval Map.empty (bor BFalse BFalse) == False

testXOR :: Bool
testXOR =
  eval Map.empty (bxor BTrue BTrue) == False &&
  eval Map.empty (bxor BTrue BFalse) == True &&
  eval Map.empty (bxor BFalse BTrue) == True &&
  eval Map.empty (bxor BFalse BFalse) == False

testHalfAdderSum :: Bool
testHalfAdderSum =
  let (s, _) = halfAdder BFalse BFalse in eval Map.empty s == False &&
  let (s, _) = halfAdder BTrue BFalse in eval Map.empty s == True &&
  let (s, _) = halfAdder BFalse BTrue in eval Map.empty s == True &&
  let (s, _) = halfAdder BTrue BTrue in eval Map.empty s == False

testHalfAdderCarry :: Bool
testHalfAdderCarry =
  let (_, c) = halfAdder BFalse BFalse in eval Map.empty c == False &&
  let (_, c) = halfAdder BTrue BFalse in eval Map.empty c == False &&
  let (_, c) = halfAdder BFalse BTrue in eval Map.empty c == False &&
  let (_, c) = halfAdder BTrue BTrue in eval Map.empty c == True

testBitVecAdd :: Bool
testBitVecAdd =
  let a = bvConst 4 3
      b = bvConst 4 5
      result = bvAdd a b
      check = bvEq result (bvConst 4 8)
  in eval Map.empty check == True

testTseitin :: Bool
testTseitin =
  let expr = band (BVar 1) (BVar 2)
      cnf = tseitin expr
  in cnfNumVars cnf > 0 && not (null (cnfClauses cnf))

testDPLLSat :: Bool
testDPLLSat =
  let cnf = CNF [[Pos 1, Pos 2], [Neg 1, Pos 2]] 2
  in case dpll cnf of
       SAT _ -> True
       UNSAT -> False

testDPLLUnsat :: Bool
testDPLLUnsat =
  let cnf = CNF [[Pos 1], [Neg 1]] 1
  in case dpll cnf of
       UNSAT -> True
       SAT _ -> False

testUnitProp :: Bool
testUnitProp =
  let clauses = [[Pos 1], [Neg 1, Pos 2]]
  in case unitPropagate Map.empty clauses of
       Propagated asgn _ -> Map.lookup 1 asgn == Just True
       Conflict -> False

testProofValid :: Bool
testProofValid =
  let cert = proveValidity "test_true" BTrue
  in proofConclusion cert == Valid

testCheckerAccepts :: Bool
testCheckerAccepts =
  let cert = proveValidity "trivial" BTrue
  in case checkCertificate cert of
       Verified _ -> True
       Rejected _ _ -> False

testCheckerRejects :: Bool
testCheckerRejects =
  let cert = (emptyProof "bad" Valid)
              { proofSteps = [Resolution [Pos 1] [Pos 2] [Neg 3] 99] }
  in case checkCertificate cert of
       Rejected _ _ -> True
       Verified _ -> False

testParse :: Bool
testParse =
  case parseModule "test" "def not(x) = (x | x);\nassert not(true);" of
    Right _ -> True
    Left _ -> False

testElaborate :: Bool
testElaborate =
  case parseModule "test" "def id(x) = x;\nassert id(true);" of
    Right m -> case elaborate m of
                 Right _ -> True
                 Left _ -> False
    Left _ -> False

testNANDNormal :: Bool
testNANDNormal =
  isNANDNormal (BNand (BVar 1) (BVar 2)) &&
  not (isNANDNormal BTrue) &&
  not (isNANDNormal BFalse)

testDeMorgan :: Bool
testDeMorgan =
  let env tt = Map.fromList [(1, fst tt), (2, snd tt)]
      cases = [(True, True), (True, False), (False, True), (False, False)]
      lhs e = bnot (band (BVar 1) (BVar 2))
      rhs e = bor (bnot (BVar 1)) (bnot (BVar 2))
  in all (\c -> eval (env c) (lhs c) == eval (env c) (rhs c)) cases

testMux :: Bool
testMux =
  let mux sel a b = bor (band (bnot sel) a) (band sel b)
      cases = [ (False, False, False, False)
              , (False, False, True, False)
              , (False, True, False, True)
              , (False, True, True, True)
              , (True, False, False, False)
              , (True, False, True, True)
              , (True, True, False, False)
              , (True, True, True, True)
              ]
      check (s, a, b, expected) =
        let se = if s then BTrue else BFalse
            ae = if a then BTrue else BFalse
            be = if b then BTrue else BFalse
        in eval Map.empty (mux se ae be) == expected
  in all check cases