import Prelude hiding ( not )
import qualified Prelude

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

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

-- | command line arguments: c_1 .. c_k n
-- program prints graph g that proves
-- R(c_1, .., c_k) > n

main :: IO ()
main = do
    argv <- fmap ( map read ) getArgs
    let cs = init argv
        n  = last argv
    Just a <- solve $ ramsey cs n
    print a

type Graph = Relation Int Int

ramsey cs n = do
    cols <- sequence $ replicate (length cs) $ do
        r <- relation ((1,1),(n,n))
        monadic assert [ symmetric r ]
        monadic assert [ irreflexive r ]
        return r
    circular_colouring ( n `div` length cs ) cols
    each_edge_is_coloured n cols
    forM ( zip cs cols ) no_monochromatic_clique
    return $ do
        ds <- mapM decode cols
        return $ do
            i <- range ((1,1),(n,n))
            let c = length $ takeWhile Prelude.not $ do d <- ds ; return $ d A.! i
            return ( i, c )

circular_colouring period cols = sequence_ $ do
    (col, col') <- zip cols $ rotate 1 cols
    x @ (p,q) <- indices col
    let y = (p+period,q+period)
    guard $ inRange ( bounds col ) y
    return $ do
        assert [ not $ col ! x, col' ! y ]

rotate :: Int -> [a] -> [a]
rotate k xs = 
    let ( pre, post ) = splitAt k xs
    in  post ++ pre

each_edge_is_coloured n cols = sequence_ $ do
    (p,q) <- range ((1,1),(n,n))
    guard $ p < q
    return $ assert $ do 
            col <- cols
            return $ col ! (p,q)

no_monochromatic_clique (c, col) = sequence_ $ do
    let ((lo,_),(hi,_)) = bounds col
    xs <- ordered_sublists c [lo .. hi]
    return $ assert $ do
        x : ys <- tails xs
        y <- ys
        return $ not $ col!(x,y)

ordered_sublists :: Int -> [a] -> [[a]]
ordered_sublists 0 xs = return []
ordered_sublists k xs = do
    ( pre, this : post ) <- splits xs
    that <- ordered_sublists (k-1) $ post
    return $ this : that

splits :: [a] -> [ ([a],[a]) ]
splits xs = zip ( inits xs ) ( tails xs )