Solve.hs

module GraphColoring.Solve where

import qualified Math.SetCover.Exact as ESC

import Control.Monad (guard)
import Control.Applicative (liftA2)

import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.List as List
import Data.Set (Set)
import Data.Ord.HT (comparing)
import Data.Maybe (catMaybes)
import Text.Printf (printf)

{- $setup
>>> import qualified GraphColoring.CutEnumSet as CutEnumSet
>>> import qualified Math.SetCover.Exact as ESC
>>>
>>> import qualified Data.Set as Set
>>> import qualified Data.List.Key as Key
>>> import Data.Either.HT (maybeRight)
>>> import Data.Maybe (catMaybes, mapMaybe)
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.QuickCheck ((===))
>>>
>>> solve :: ([Char], [Edge Int]) -> [[(Int, Char)]]
>>> solve =
>>>    map (Key.sort fst . catMaybes) . ESC.search .
>>>    ESC.compressState Set.toList . initialize assigns
-}


type Edge node = (node,node)

data X node color =
      Node node | Edge node node | ColorEdge node node color |
      EdgeColor node node color
         deriving (Eq, Ord, Show)

type Label node color = Maybe (node, color)
type ExtLabel node color = Either (color,color) (node, color)
type Assign node color = ESC.Assign (Label node color) (Set (X node color))


{- |
For every node, block all but one color at outgoing edges.

Specify every edge only once, freely choose a direction.
Strange things might happend if you specify an edge in both directions.
-}
assignNodeColors ::
   (Ord color, Ord node) =>
   ((node, color) -> lab) ->
   [color] -> [Edge node] -> [ESC.Assign lab (Set (X node color))]
assignNodeColors lab colors edges =
   liftA2
      (\color (node,connected) ->
         ESC.assign (lab (node,color))
            (Set.fromList $
               Node node :
               liftA2 (ColorEdge node)
                  (Set.toList connected)
                  (List.delete color colors)))
      colors
      (Map.toList $ Map.fromListWith Set.union $
       edges >>=
         \(from,to) -> [(from, Set.singleton to), (to, Set.singleton from)])

{- FixMe:
Solver always succeeds when only one color is allowed.
When there is only one color, we do not generate edges, at all,
because we only generate edges with different colors at either ends.
-}
assigns ::
   (Ord color, Ord node) => [color] -> [Edge node] -> [Assign node color]
assigns colors edges =
   assignNodeColors Just colors edges
   ++
   (do
      (from,to) <- edges
      fromColor <- colors
      toColor <- colors
      guard $ fromColor /= toColor
      return $ ESC.assign Nothing $
         Set.fromList [ColorEdge from to fromColor, ColorEdge to from toColor])

{- |
Register every used edge.
This is redundant but may speedup search.
Actually, it slows down the search.
-}
assignsReserveEdge ::
   (Ord color, Ord node) => [color] -> [Edge node] -> [Assign node color]
assignsReserveEdge colors edges =
   assignNodeColors Just colors edges
   ++
   (do
      (from,to) <- edges
      fromColor <- colors
      toColor <- colors
      guard $ fromColor /= toColor
      return $ ESC.assign Nothing $
         Set.fromList
            [ColorEdge from to fromColor,
             ColorEdge to from toColor,
             Edge from to])

{- |
For every edge, 'assigns' needs a number of sets
proportional to the square of the number of colors.
'assignsSplitEdge' reduces this to a linear number.
Unfortunately solution becomes slower.
-}
assignsSplitEdge ::
   (Ord color, Ord node) => [color] -> [Edge node] -> [Assign node color]
assignsSplitEdge colors edges =
   assignNodeColors Just colors edges
   ++
   (do
      (from,to) <- edges
      color <- colors
      [ESC.assign Nothing $ Set.fromList [EdgeColor from to color],
       ESC.assign Nothing $
         Set.fromList
            [ColorEdge from to color, EdgeColor from to color, Edge from to],
       ESC.assign Nothing $
         Set.fromList
            [ColorEdge to from color, EdgeColor from to color, Edge to from]
         ])


{- |
Same rules as 'assign' but also gives labels for edge colors.
-}
assignsEdgeColor ::
   (Ord color, Ord node) =>
   [color] -> [Edge node] ->
   [ESC.Assign (ExtLabel node color) (Set (X node color))]
assignsEdgeColor colors edges =
   assignNodeColors Right colors edges
   ++
   (do
      (from,to) <- edges
      fromColor <- colors
      toColor <- colors
      guard $ fromColor /= toColor
      return $ ESC.assign (Left (fromColor, toColor)) $
         Set.fromList [ColorEdge from to fromColor, ColorEdge to from toColor])



{- |
>>> solve example0
[[(0,'a'),(1,'b'),(2,'a'),(3,'b'),(4,'a'),(5,'b')]]
-}
example0 :: ([Char], [Edge Int])
example0 = (['a'..'b'], [(0,1),(1,2),(2,3),(3,4),(4,5)])

{- |
>>> minimum $ solve example1
[(0,'a'),(1,'b'),(2,'a'),(3,'b'),(4,'c')]
-}
example1 :: ([Char], [Edge Int])
example1 = (['a'..'c'], [(0,1),(1,2),(2,3),(3,4),(4,0)])

{- |
>>> solve example2
[[(0,'a'),(1,'b'),(2,'c')]]
-}
example2 :: ([Char], [Edge Int])
example2 = (['a'..'c'], [(0,1),(1,2),(0,2)])

{- |
>>> solve example3
[[(0,'a'),(1,'b'),(2,'c'),(3,'b')]]
-}
example3 :: ([Char], [Edge Int])
example3 = (['a'..'c'], [(0,1),(1,2),(2,3),(3,0),(0,2)])

{- |
>>> solve example4
[]
-}
example4 :: ([Char], [Edge Int])
example4 = (['a'..'c'], [(0,1),(1,2),(2,3),(3,0),(0,2),(1,3)])

-- FixMe: fix cases n<2 or k<2
{- |
>>> length $ ESC.search $ ESC.compressState Set.toList $ ESC.initState $ assigns ['a'..'e'] (complete 5)
120

prop> :{
   QC.forAll (QC.choose (2,5)) $ \n ->
   QC.forAll (QC.choose (2,7)) $ \k ->

   (length $ ESC.search $ ESC.compressState Set.toList $
      ESC.initState $ assigns (take k ['a'..]) (complete n))
   ===
   product (take n [k,(k-1)..])
:}

prop> :{
   QC.forAll (QC.choose (2,5)) $ \n ->
   QC.forAll (QC.choose (2,7)) $ \k ->

   (length $ CutEnumSet.search $ CutEnumSet.initState $
      ESC.intSetFromSetAssigns $ assignsEdgeColor (take k ['a'..]) (complete n))
   ===
   if n<=k then 1 else 0
:}
-}
complete :: Int -> [Edge Int]
complete n = do from <- [1..n]; to <- [from+1 .. n]; [(from,to)]

{- |
>>> :{
   map (Key.sort fst . mapMaybe maybeRight) $
   CutEnumSet.search $ CutEnumSet.initState $
   ESC.intSetFromSetAssigns $ uncurry assignsEdgeColor example5
:}
[[(1,'a'),(2,'b'),(3,'c'),(4,'d'),(5,'e'),(6,'f'),(7,'g'),(8,'h'),(9,'i'),(10,'j')]]
-}
example5 :: ([Char], [Edge Int])
example5 = (['a'..'j'], complete 10)



{- |
Choose first two colors for the first two connected nodes,
without loss of generality.
This accelerates proof of impossibility.
However, for the following colors still plain color permutations are tried.
We could exclude them by modifying the search algorithm
such that at most one unused color is tried in every search step.
-}
initialize ::
   (Eq node, Eq color, ESC.Set set) =>
   ([color] -> [Edge node] -> [ESC.Assign (Label node color) set]) ->
   ([color], [Edge node]) -> ESC.State (Label node color) set
initialize assigner (colors,edges) =
   let asns = assigner colors edges in
   case (colors,edges) of
      (c0:c1:_, (from,to):_) ->
         foldl (flip ESC.updateState) (ESC.initState asns) $
         filter
            (\asn ->
               ESC.label asn == Just (from,c0) ||
               ESC.label asn == Just (to,c1))
            asns
      _ -> ESC.initState asns

{- |
The algorithm needs a fixed number of colors to choose from.
In some applications the number of colors is not given, but shall be minimal.
In this case you might start with as many colors as nodes,
and scan the results for solutions with less used colors.
You can then restart the algorithm with a smaller reservoir of colors.
You may also perform a bisection on the number of colors
or run solvers for multiple color sets in parallel.
I guess, for large color sets you will get a positive result quickly
and for small color sets you will get a negative answer quickly.

Finding a solution usually happens fast.
Proof of impossibility requires much more time.
For a real application you might consider setting a timeout
or trying to solve with different sets of colors in parallel.
-}
run :: (ESC.Set set) => ESC.State (Maybe (Int, Char)) set -> IO ()
run state =
   case ESC.search state of
      [] -> putStrLn "coloring impossible"
      solution:_ ->
         mapM_ (\(node,color) -> printf "Node %d: %c\n" node color) $
         List.sortBy (comparing fst) $ catMaybes solution