{-# LANGUAGE CPP #-}
module Main where

import Prelude hiding (not, (&&), and, or)

import Ersatz
import Ersatz.Relation
import Ersatz.Counting

import qualified Data.Array as A
import Data.Array (Array, Ix)

import Control.Monad.State.Lazy (StateT)

import Data.List (tails)

import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy.Char8 as C
import Data.Functor.Identity
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup(..))
#endif

import Options.Applicative
import Text.Printf (printf)


data Coloring = Coloring Int Int Int

coloring :: Parser Coloring
coloring = Coloring
  <$> argument auto ( metavar "n" <> showDefault <> value 11 <> help "Number of nodes of graph G" )
  <*> argument auto ( metavar "c" <> showDefault <> value 3 <> help "G is not c-colorable" )
  <*> argument auto ( metavar "k" <> showDefault <> value 3 <> help "G has no complete subgraph with k nodes" )

-- | default: Grötzsch graph: not 3-colorable and K3-free (11 nodes)
main :: IO ()
main = execParser options >>= run
  where options = info (coloring <**> helper) fullDesc

run :: Coloring -> IO ()
run (Coloring n c k) = do
  printf "n = %d, c = %d, k = %d\n" n c k
  formulaSize $ problem n c k
  solve $ problem n c k


solve :: StateT QSAT IO (Relation Int Int) -> IO ()
solve p = do
  result <- solveWith depqbf p
  case result of
    (Satisfied, Just r) -> mapM_ putStrLn [table r, show $ edgesA r]
    _                   -> putStrLn "unsat"

-- | @problem n c k@ generates a QBF problem that encodes a graph with @n@ nodes,
-- which is not @c@-colorable and does not contain a complete subgraph with @k@ nodes.
problem :: Monad a => Int -> Int -> Int -> StateT QSAT a (Relation Int Int)
problem n c k = do
  r <- symmetric_relation ((0,0),(n-1,n-1))
  col <- universally_quantified_relation ((0,0),(n-1,c-1))
  assert $ and [
      irreflexive r
    , not $ has_k k r
    , is_coloring col ==> not $ proper col r
    ]
  return r


universally_quantified_relation :: (Ix a, Ix b, MonadQSAT s m)
  => ((a,b),(a,b)) -> m (Relation a b)
universally_quantified_relation bnd = do
  pairs <- sequence $ do
    p <- A.range bnd
    return $ do
      x <- forall_
      return (p,x)
  return $ build bnd pairs

-- | @has_k n r@ encodes the constraint that @r@ has a complete subgraph with @n@ nodes.
has_k :: Ix a => Int -> Relation a a -> Bit
has_k n r = or $ do
  xss <- select n $ universe r
  return $ and $ do
    (x:xs) <- tails xss
    y <- xs
    return $ r!(x,y)

-- | select 2 [1..4] = [[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]]
select :: Int -> [a] -> [[a]]
select 0 _ = [[]]
select _ [] = []
select k (x:xs) = map (x:) (select (k-1) xs) ++ select k xs

-- | Given a relation r with domain a and codomain b, check if r matches every element in a
-- to exactly one element in b.
is_coloring :: Ix a => Relation a a -> Bit
is_coloring c = and $ do
  i <- domain c
  return $ exactly 1 $ do
    j <- codomain c
    return $ c!(i,j)

-- | @proper c r@ encodes the constraint that @c@ is a proper coloring for @r@.
proper :: Ix a => Relation a a -> Relation a a -> Bit
proper col r = and $ do
  (p,q) <- indices r
  j <- codomain col
  return $ r!(p,q) ==> not $ col!(p,j) && col!(q,j)


edgesA :: (Ix a, Ix b) => Array (a,b) Bool -> [(a,b)]
edgesA a = [ i | (i, b) <- A.assocs a, b == True]

formulaSize :: StateT QSAT Identity a -> IO ()
formulaSize p = mapM_ C.putStrLn $ take 2 $ B.split 10 $ qdimacsQSAT p