import Prelude hiding ( not, or, and )

import Satchmo.Relation
import Satchmo.Code
import Satchmo.Boolean
import Satchmo.Counting
import Satchmo.Solve

import Data.List ( inits, tails )
import qualified Data.Array as A
import System.Environment
import Control.Monad ( guard, forM_ )

-- | command line arguments: c n
-- program looks for sum-free c-colouring of [1 .. n]

main :: IO ()
main = do
    argv <- getArgs
    let [ c, n ] = map read argv
    Just a <- solve $ schur c n
    putStrLn $ table a
    print $ do
        o <- [ 1 .. c ]
        return ( o, length $ do i <- [ 1 .. n ]; guard $ a A.! (i,o) )

schur c n = do
    col <- relation ((1,1),(n,c))
    each_number_coloured col
    sum_free_colouring col
    return $ decode col

periodic p col = sequence_ $ do
    let ((1,1),(n,c)) = bounds col
    x <- [ 1 .. n ]
    let y = x + p
    guard $ y <= n
    o <- [ 1 .. c ]
    let p = 1 + o `mod` c
    return $ assert [ not $ col!(x,o), col!(y,p) ]

each_number_coloured col = sequence_ $ do
    let ((1,1),(n,c)) = bounds col
    x <- [ 1 .. n ]
    return $ assert $ do o <- [1 .. c]; return $ col!(x,o)

sum_free_colouring col = sequence_ $ do
    let ((1,1),(n,c)) = bounds col
    x <- [ 1 .. n ]
    y <- [ x .. n ]
    let z = (x + y) `mod` (n+1)
    guard $ z <= n
    guard $ 1 <= z
    o <- [1 .. c]
    return $ assert $ do 
        p <- [ x, y, z ]
        return $ not $ col!(p,o)

evenly_distributed col = do
    let ((1,1),(n,c)) = bounds col
        d = n `div` c
    forM_ [ 1 .. c ] $ \ o -> do
        a <- atleast d $ do i <- [ 1 .. n ] ; return $ col!(i,o)
        assert [a]