content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# OPTIONS_GHC -Wall #-}
  module Main where
  
  import Control.Monad
  import Data.Array.IArray
  import qualified Data.ByteString.Lazy.Char8 as BL
  import Data.Char
  import Data.Default.Class
  import Data.List (intersperse)
  import Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import Data.Maybe
  import qualified Data.PseudoBoolean as PB
  import Data.Set (Set)
  import qualified Data.Set as Set
  import System.Console.GetOpt
  import System.Environment
  import System.Exit
  import System.IO
  import Text.Parsec hiding (try)
  import qualified Text.Parsec.ByteString.Lazy as ParsecBL
  import qualified ToySolver.FileFormat as FF
  import qualified ToySolver.SAT as SAT
  import qualified ToySolver.SAT.PBO as PBO
  import qualified ToySolver.SAT.Store.PB as PBStore
  
  data Problem
    = Problem
    { probSize :: (Int,Int,Int)
    , probIsMultiLayer :: Bool
    , probLineNum :: Int
    , probLines :: [(Number, Cell, Cell)]
    , probVias :: [(Via, [Cell])]
    }
    deriving (Show)
  
  type Number = Int
  type Cell = (Int,Int,Int)
  type Edge = (Cell, Cell)
  type Via = String
  
  parser ::  Stream s m Char => ParsecT s u m Problem
  parser = do
    optional $ string "\xEF\xBB\xBF" -- BOM in UTF-8
    spaces
    _ <- string "SIZE"
    spaces
    w <- num
    _ <- char 'X'
    h <- num
    dd <- optionMaybe $ char 'X' >> num
    let (d,isMultiLayer) =
          case dd of
            Nothing -> (1, False)
            Just d -> (d, True)
    _ <- endOfLine
  
    let cell = do
          _ <- char '('
          x <- num
          _ <- char ','
          y <- num
          z <-
            if isMultiLayer then do
              _ <- char ','
              num
            else
              return 1
          _ <- char ')'
          spaces
          return (x,y,z)
  
    _ <- string "LINE_NUM"
    spaces
    n <- num
    _ <- endOfLine
  
    xs <- many $ do
      _ <- string "LINE#"
      i <- num
      spaces
      src <- cell
      optional $ char '-' >> spaces
      dst <- cell
      return (i,src,dst)
  
    ys <- many $ do
      _ <- string "VIA#"
      v <- many letter
      spaces
      cs <- cell `sepBy` optional (char '-' >> spaces)
      return (v,cs)
  
    return $ Problem (w,h,d) (isJust dd) n xs ys
  
    where
      num = read <$> many digit
  
  type Encoded = (Array Cell (Array Number SAT.Var), Map Edge SAT.Var)
  
  encode :: SAT.AddPBLin m enc => enc -> Options -> Problem -> m Encoded
  encode enc opt Problem{ probSize = (w,h,d), probLineNum = n, probLines = ls, probVias = vias } = do
    let bnd = ((0,0,1), (w-1,h-1,d))
        cells = range bnd
        edges = [(a,b) | a@(x,y,z) <- cells, b <- [(x+1,y,z),(x,y+1,z)], inRange bnd b]
        adjacentEdges a@(x,y,z) =
          [(a,b) | b <- [(x+1,y,z),(x,y+1,z)], inRange bnd b] ++
          [(b,a) | b <- [(x-1,y,z),(x,y-1,z)], inRange bnd b]
        ks = [1..n]
  
    -- ビアへの数字の割当を表す変数 (0 は数字なしを表す)
    viaVs <- liftM Map.fromList $ forM vias $ \(via, _) -> do
      let r = if optAssumeNoBlank opt
              then (1,n)
              else (0,n)
      zs <- liftM (array r) $ forM (range r) $ \k -> do
        v <- SAT.newVar enc
        return (k,v)
      return (via, zs)
    let viaPos = Map.fromList [(a,via) | (via,as) <- vias, a <- as]
  
    -- 各マスへの数字の割り当てを表す変数 (0 は数字なしを表す)
    vs <- liftM (array bnd) $ forM cells $ \a -> do
      case Map.lookup a viaPos of
        Just via -> return (a, viaVs Map.! via)
        Nothing -> do
          let r = if optAssumeNoBlank opt
                  then (1,n)
                  else (0,n)
          zs <- liftM (array r) $ forM (range r) $ \k -> do
            v <- SAT.newVar enc
            return (k,v)
          return (a, zs)
    -- 各辺の有無を表す変数
    evs <- liftM Map.fromList $ forM edges $ \e -> do
      v <- SAT.newVar enc
      return (e,v)
  
    -- 初期数字
    let m0 = Map.fromList [(c,k) | (k,src,dst) <- ls, c <- [src,dst]]
    forM_ (Map.toList m0) $ \(c,k) -> do
      SAT.addClause enc [vs!c!k]
  
    -- 各マスには高々ひとつしか数字が入らない
    forM_ (range bnd) $ \a -> do
      SAT.addExactly enc [v | v <- elems (vs!a)] 1
    forM_ (Map.toList evs) $ \((a,b),v) ->
      forM_ ks $ \k -> do
        -- 辺で連結されるマスは同じ数字
        SAT.addClause enc [-v, -(vs!a!k), vs!b!k]
        SAT.addClause enc [-v, -(vs!b!k), vs!a!k]
        -- 連結されない隣接マスは違う数字
        when (optAssumeNoDetour opt) $
          SAT.addClause enc [v, -(vs!a!k), -(vs!b!k)]
  
    forM_ (range bnd) $ \a -> do
      if a `Map.member` m0 then do
        -- 初期数字の入っているマスの次数は1
        SAT.addExactly enc [evs Map.! e | e <- adjacentEdges a] 1
      else if a `Map.member` viaPos then do
        -- ビアの次数は高々1 (ADC2016ルールでも中間層は0になるので注意)
        SAT.addAtMost enc [evs Map.! e | e <- adjacentEdges a] 1
      else do
        -- 初期数字の入っていないマスの次数は0か2
        if optAssumeNoBlank opt then do
          SAT.addExactly enc [evs Map.! e | e <- adjacentEdges a] 2
        else do
          let v = vs!a!0
          -- v → deg(a)=0
          -- SAT.addPBAtMostSoft enc v [(1, evs Map.! e) | e <- adjacentEdges a] 0
          forM_ (adjacentEdges a) $ \e -> SAT.addClause enc [-v, -(evs Map.! e)]
          -- ¬v → deg(a)=2
          SAT.addPBExactlySoft enc (-v) [(1, evs Map.! e) | e <- adjacentEdges a] 2
          -- ¬v → deg(a)>0
          -- SAT.addClause enc $ v : [evs Map.! e | e <- adjacentEdges a]
  
    -- コの字の禁止
    when (optAssumeNoDetour opt) $ do
      forM_ [1..d] $ \z -> do
        forM_ [0..w-2] $ \x -> do
          forM_ [0..h-2] $ \y -> do
            let a = (x, y, z)
                b = (x+1, y, z)
                c = (x, y+1, z)
                d = (x+1, y+1, z)
            SAT.addAtMost enc [evs Map.! e | e <- [(a,b),(a,c),(b,d),(c,d)]] 2
  
    -- https://kaigi.org/jsai/webprogram/2016/pdf/67.pdf の追加成約
    when (optJSAI2016 opt) $ do
      let bs = Set.fromList [a | a <- range bnd, a `Map.notMember` m0, a `Map.notMember` viaPos]
      forM_ (range bnd) $ \a@(x,y,z) -> do
        let a_n = (x,y-1,z)
            a_s = (x,y+1,z)
            a_w = (x-1,y,z)
            a_e = (x+1,y,z)
            a_nw = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x-1,y-1,z)) a
            a_ne = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x+1,y-1,z)) a
            a_sw = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x-1,y+1,z)) a
            a_se = Set.fromList $ takeWhile (inRange bnd) $ tail $ iterate (\(x,y,z) -> (x+1,y+1,z)) a
        when (inRange bnd a_n && inRange bnd a_w && a_nw `Set.isSubsetOf` bs) $ do
          SAT.addAtMost enc [evs Map.! (a_w,a), evs Map.! (a_n,a)] 1
        when (inRange bnd a_n && inRange bnd a_e && a_ne `Set.isSubsetOf` bs) $ do
          SAT.addAtMost enc [evs Map.! (a,a_e), evs Map.! (a_n,a)] 1
        when (inRange bnd a_s && inRange bnd a_w && a_sw `Set.isSubsetOf` bs) $ do
          SAT.addAtMost enc [evs Map.! (a_w,a), evs Map.! (a,a_s)] 1
        when (inRange bnd a_s && inRange bnd a_e && a_se `Set.isSubsetOf` bs) $ do
          SAT.addAtMost enc [evs Map.! (a,a_e), evs Map.! (a,a_s)] 1
  
    forM_ ls $ \(k, (_,_,z1), (_,_,z2)) -> do
      let vsk = [vs!k | (_, vs) <- Map.toList viaVs]
      if z1 == z2 then do
        -- ADC2016では(層をまたぐLINE数)=(VIA数)なので、他のLINEはVIAを使えない
        when (optADC2016 opt) $ SAT.addAtMost enc vsk 0
      else do
        -- 異なる盤面にある数字はビアを必ず使う
        SAT.addClause enc vsk
        -- ADC2016では(層をまたぐLINE数)=(VIA数)
        when (optADC2016 opt) $ SAT.addAtMost enc vsk 1
  
    return (vs, evs)
  
  encodeObj :: SAT.AddPBLin m enc => enc -> Options -> Problem -> Encoded -> m SAT.PBLinSum
  encodeObj enc opt Problem{ probSize = (w,h,d) } (cells,edges) = do
    let (o1, o2) = fromJust (optOptimize opt)
    obj1 <-
      if not o1 then do
        return []
      else if optAssumeNoBlank opt then do
        v <- SAT.newVar enc
        SAT.addClause enc [v]
        return [(fromIntegral (w*h*d), v)]
      else
        return [(1, -(vs!0)) | vs <- elems cells]
    obj2 <-
      if not o2 then do
        return []
      else do
        forM (range (bounds cells)) $ \a@(x,y,z) -> do
          let w = ((x-1,y,z),a)
              e = (a,(x+1,y,z))
              n = ((x,y-1,z),a)
              s = (a,(x,y+1,z))
          v <- SAT.newVar enc
          forM_ [e,w] $ \e1 -> do
            case Map.lookup e1 edges of
              Nothing -> return ()
              Just v1 -> do
                forM_ [n,s] $ \e2 -> do
                  case Map.lookup e2 edges of
                    Nothing -> return ()
                    Just v2 -> SAT.addClause enc [-v1, -v2, v]
          return (1,v)
    return $ obj1 ++ obj2
  
  type Solution = (Map Cell Number, Set Edge)
  
  decode :: Problem -> Encoded -> SAT.Model -> Solution
  decode prob (vs, evs) m = (solCells, solEdges)
    where
      solCells = Map.fromList $ do
        a <- range (bounds vs)
        guard $ a `Set.member` usedCells
        case [k | (k,v) <- assocs (vs!a), SAT.evalLit m v] of
          k:_ -> return (a,k)
          [] -> mzero
      solEdges = Set.fromList [e | e@(a,_) <- edges, a `Set.member` usedCells]
  
      edges = [e | (e,ev) <- Map.toList evs, SAT.evalLit m ev]
      adjacents = Map.fromListWith Set.union $ concat $ [[(a, Set.singleton b), (b, Set.singleton a)] | (a,b) <- edges]
      usedVias = Set.fromList [(x,y,z) | ((x,y),zs) <- Map.toList cs, z <- [Set.findMin zs .. Set.findMax zs]]
        where
          cs = Map.fromList [((x,y), Set.singleton z) | (_,as) <- probVias prob, a@(x,y,z) <- as, a `Map.member` adjacents]
      usedCells = Set.union usedVias cs
        where
          cs = Set.unions [g (Set.fromList [a,b]) Set.empty | (_k,a,b) <- probLines prob]
          g xs visited
            | Set.null xs = visited
            | otherwise = g (Set.unions [Map.findWithDefault Set.empty x adjacents Set.\\ visited | x <- Set.toList xs]) (Set.union xs visited)
  
  evalObj :: Options -> Problem -> Solution -> Integer
  evalObj opt Problem{ probSize = (w,h,d) } (cells,edges) = obj1 + obj2
    where
      (o1, o2) = fromJust (optOptimize opt)
      bnd = ((0,0,1),(w-1,h-1,d))
      obj1
        | not o1 = 0
        | otherwise = fromIntegral $ Map.size cells
      obj2
        | not o2 = 0
        | otherwise = sum $ do
            a@(x,y,z) <- range bnd
            let w = ((x-1,y,z),a)
                e = (a,(x+1,y,z))
                n = ((x,y-1,z),a)
                s = (a,(x,y+1,z))
            if null [() | e1 <- [e,w], e1 `Set.member` edges, e2 <- [n,s], e2 `Set.member` edges] then
              return 0
            else
              return 1
  
  createSolver :: Options -> Problem -> IO (IO (Maybe Solution))
  createSolver opt prob = do
    solver <- SAT.newSolver
    SAT.setLogger solver $ hPutStrLn stderr
    encoded <- encode solver opt prob
    unless (optAssumeNoBlank opt) $ do
      forM_ (elems (fst encoded)) $ \xs -> do
        SAT.setVarPolarity solver (xs!0) False
    let m = do
          ret <- SAT.solve solver
          if ret then do
            m <- SAT.getModel solver
            let sol = decode prob encoded m
            SAT.addClause solver $ blockingClause prob encoded sol
            return $ Just sol
          else
            return Nothing
    return m
  
  printSolution :: Problem -> Solution -> IO ()
  printSolution prob (cells, _) = do
    let (w,h,d) = probSize prob
    forM_ [1 .. d] $ \z -> do
      when (probIsMultiLayer prob) $ do
        unless (z == 1) $ putStrLn ""
        putStrLn $ "LAYER " ++ show z
      forM_ [0 .. h-1] $ \y -> do
        putStrLn $ concat $ intersperse ","
         [ case Map.lookup (x,y,z) cells of
             Nothing -> replicate width '0'
             Just k -> replicate (width - length (show k)) '0' ++ show k
         | x <- [0 .. w-1]
         ]
    where
      width = length $ show (probLineNum prob)
  
  blockingClause :: Problem -> Encoded -> Solution -> SAT.Clause
  blockingClause _prob (_,edgesEnc) (_,edges) = [- (edgesEnc Map.! e) | e <- Set.toList edges]
  
  data Options
    = Options
    { optOptimize :: Maybe Obj
    , optAssumeNoBlank :: Bool
    , optAssumeNoDetour :: Bool
    , optJSAI2016 :: Bool
    , optADC2016 :: Bool
    , optNumSolutions :: Int
    }
  
  instance Default Options where
    def =
      Options
      { optOptimize = Nothing
      , optAssumeNoBlank = False
      , optAssumeNoDetour = False
      , optJSAI2016 = False
      , optADC2016 = False
      , optNumSolutions = 1
      }
  
  type Obj = (Bool,Bool)
  
  options :: [OptDescr (Options -> Options)]
  options =
      [ Option [] ["optimize"]
          (OptArg (\arg opt -> opt{ optOptimize = Just (maybe (True,True) readObj arg) }) "<str>")
          "Specify objective functions (-n option will be ignored): both (default), length, corners"
      , Option [] ["no-blank"]
          (NoArg (\opt -> opt{ optAssumeNoBlank = True }))
          "Assume no blank cells in solution"
      , Option [] ["no-detour"]
          (NoArg (\opt -> opt{ optAssumeNoDetour = True }))
          "Assume no detour in solution"
      , Option [] ["jsai2016"]
          (NoArg (\opt -> opt{ optJSAI2016 = True }))
          "Add constraints of JSATI2016 paper of Tatsuya Seko et.al."
      , Option [] ["adc2016"]
          (NoArg (\opt -> opt{ optADC2016 = True }))
          "Add constraints specific to ADC2016 rule"
      , Option ['n'] []
          (ReqArg (\val opt -> opt{ optNumSolutions = read val }) "<int>")
          "Maximal number of solutions to enumerate, or negative value for all solutions (default: 1)"
      ]
    where
      readObj = f . map toLower
        where
          f "both" = (True, True)
          f "length" = (True, False)
          f "corners" = (False, True)
          f _ = error "unknown objective"
  
  header :: String
  header = unlines
    [ "Usage:"
    , "  numberlink [OPTION]... problem.txt"
    , "  numberlink [OPTION]... problem.txt encoded.opb"
    , ""
    , "Options:"
    ]
  
  showHelp :: Handle -> IO ()
  showHelp h = hPutStrLn h (usageInfo header options)
  
  main :: IO ()
  main = do
    args <- getArgs
    case getOpt Permute options args of
      (_,_,errs@(_:_)) -> do
        mapM_ putStrLn errs
        exitFailure
      (o,args2,[]) -> do
        let opt = foldl (flip id) def o
        case args2 of
          [fname] -> do
            r <- ParsecBL.parseFromFile parser fname
            case r of
              Left err -> error (show err)
              Right prob@Problem{ probSize = (w,h,d) } -> do
                putStrLn $ "SIZE " ++ show w ++ "X" ++ show h ++
                             (if probIsMultiLayer prob then "X" ++ show d else "")
                case optOptimize opt of
                  Nothing -> do
                    solve <- createSolver opt prob
                    let loop n
                          | n == 0 = return ()
                          | otherwise = do
                              m <- solve
                              case m of
                                Nothing -> return ()
                                Just sol -> do
                                  printSolution prob sol
                                  putStrLn ""
                                  loop (if n > 0 then n - 1 else n)
                    loop (optNumSolutions opt)
                  Just _ -> do
                    solver <- SAT.newSolver
                    SAT.setLogger solver $ hPutStrLn stderr
                    encoded@(cells,_edges) <- encode solver opt prob
                    unless (optAssumeNoBlank opt) $ do
                      forM_ (elems cells) $ \xs -> do
                        SAT.setVarPolarity solver (xs!0) False
                    obj <- encodeObj solver opt prob encoded
                    pbo <- PBO.newOptimizer2 solver obj (\m -> evalObj opt prob (decode prob encoded m))
                    PBO.setLogger pbo $ hPutStrLn stderr
                    PBO.setOnUpdateBestSolution pbo $ \m val -> do
                      hPutStrLn stderr $ "# obj = " ++ show val
                      hFlush stderr
                      let sol = decode prob encoded m
                      printSolution prob sol
                      putStrLn ""
                      hFlush stdout
                    PBO.optimize pbo
  
          [fname, fname2] -> do
            r <- ParsecBL.parseFromFile parser fname
            case r of
              Left err -> error (show err)
              Right prob -> do
                store <- PBStore.newPBStore
                obj <- do
                  encoded <- encode store opt prob
                  case optOptimize opt of
                    Nothing -> return Nothing
                    Just _ -> do
                      obj <- encodeObj store opt prob encoded
                      return $ Just [(c,[v]) | (c,v) <- obj]
                opb <- PBStore.getPBFormula store
                FF.writeFile fname2 $ opb{ PB.pbObjectiveFunction = obj }
          _ -> do
            showHelp stderr
  
  sampleFile :: BL.ByteString
  sampleFile = BL.unlines
    [ "SIZE 15X15"
    , "LINE_NUM 12"
    , "LINE#1 (10,0)-(4,14)"
    , "LINE#2 (9,5)-(9,14)"
    , "LINE#3 (2,7)-(4,12)"
    , "LINE#4 (2,5)-(7,13)"
    , "LINE#5 (2,4)-(4,10)"
    , "LINE#6 (2,2)-(5,5)"
    , "LINE#7 (6,5)-(5,10)"
    , "LINE#8 (5,0)-(7,9)"
    , "LINE#9 (10,2)-(12,7)"
    , "LINE#10 (7,1)-(12,9)"
    , "LINE#11 (9,8)-(12,12)"
    , "LINE#12 (8,9)-(12,10)"
    ]
  
  sample :: Problem
  Right sample = parse parser "sample" sampleFile
  