module CountryColoring where

import EuropeBorders
import Z3.Monad
import Control.Monad
import Data.Map (Map)
import qualified Data.Map as M
import qualified Data.Traversable as T
import Data.Maybe

data Color = Red | Green | Blue | Yellow
  deriving (Show, Eq, Ord, Enum, Bounded)

type Coloring = Map Country Color

{-
Реализуйте функцию, которая с помощью Z3 вычисляет такую раскраску карты
с заданными границами между странами, чтобы никакие две соседние страны
не были бы покрашены одним цветом. Карта задаётся списками соседних стран
для каждой страны. Пример такой карты имеется в модуле EuropeBorders.
Известно, что для раскраски любой карты достаточно четырёх цветов
(https://en.wikipedia.org/wiki/Four_color_theorem).

Подсказка: цвет каждой страны в формуле для SMT-солвера -- это целая
переменная, равная 0, 1, 2 или 3, ограниченная также набором
неравенств для всех стран, граничащих с данной страной.
-}

coloring :: Map Country [Country] -> Z3 (Maybe Coloring)
coloring mapCountry = do
  nodes <- foldM (\acc x -> do
    var <- mkFreshIntVar x
    return $ M.insert var x acc
                 ) M.empty (M.keys mapCountry)
  _0 <- mkInteger (0 :: Integer)
  _3 <- mkInteger (3 :: Integer)
  assert =<< mkAnd =<< T.sequence [ mkLe _0 v | v <- M.keys nodes]
  assert =<< mkAnd =<< T.sequence [ mkLe v _3 | v <- M.keys nodes]
  assert =<< mkNot =<< mkOr =<< T.sequence [
    mkEq v1 v2 | v1 <- M.keys nodes, v2 <- M.keys nodes,
    (nodes M.! v1) `elem` (mapCountry M.! (nodes M.! v2))]
  getSolution nodes
  where
    getSolution :: Map AST String -> Z3 (Maybe Coloring)
    getSolution nodes = do
      let vals = M.keys nodes
      colors <- (fmap . fmap . fmap) toColor
        (fmap snd $ withModel $ \m -> catMaybes <$> mapM (evalInt m) vals)
      let colorMap = fmap (\cs -> zipWith (\a b -> (a, b)) [nodes M.! key | key <- vals] cs) colors
      return $ fmap M.fromList colorMap
    toColor i = M.fromList [(0, Red), (1, Green), (2, Blue), (3, Yellow)] M.! i