content: {-# LANGUAGE CPP #-}
  {-# OPTIONS_GHC -Wall #-}
  -----------------------------------------------------------------------------
  -- |
  -- Module      :  toysolver
  -- Copyright   :  (c) Masahiro Sakai 2011-2019
  -- License     :  BSD-style
  --
  -- Maintainer  :  masahiro.sakai@gmail.com
  -- Stability   :  experimental
  -- Portability :  non-portable
  --
  -----------------------------------------------------------------------------
  
  
  module Main where
  
  import Control.Monad
  import Control.Concurrent
  import Data.Array.IArray
  import Data.Char
  import Data.Default.Class
  import Data.IORef
  import Data.List (intercalate)
  import Data.Maybe
  import Data.Scientific (Scientific)
  import qualified Data.Scientific as Scientific
  import Data.String
  import qualified Data.Version as V
  import qualified Data.Set as Set
  import Data.Map (Map)
  import qualified Data.Map as Map
  import qualified Data.IntMap as IntMap
  import qualified Data.IntSet as IntSet
  import qualified Data.Traversable as T
  import Options.Applicative hiding (Const)
  import System.Exit
  import System.IO
  import Text.Printf
  import GHC.Conc (getNumProcessors)
  
  import qualified Numeric.Optimization.MIP as MIP
  import qualified Numeric.Optimization.MIP.Solution.Gurobi as GurobiSol
  
  import ToySolver.Data.OrdRel
  import ToySolver.Data.FOL.Arith as FOL
  import qualified ToySolver.Data.LA.FOL as LAFOL
  import qualified ToySolver.Data.Polynomial as P
  import qualified ToySolver.Data.AlgebraicNumber.Real as AReal
  import qualified ToySolver.Arith.OmegaTest as OmegaTest
  import qualified ToySolver.Arith.Cooper as Cooper
  import qualified ToySolver.Arith.Simplex.Textbook.MIPSolver.Simple as TextbookMIP
  import qualified ToySolver.Arith.Simplex as Simplex
  import qualified ToySolver.Arith.MIP as MIPSolver
  import qualified ToySolver.Arith.CAD as CAD
  import qualified ToySolver.Arith.ContiTraverso as ContiTraverso
  import qualified ToySolver.FileFormat as FF
  import ToySolver.Converter
  import ToySolver.SAT.Printer
  import qualified ToySolver.SAT.Types as SAT
  import ToySolver.Version
  import ToySolver.Internal.Util
  
  -- ---------------------------------------------------------------------------
  
  data Mode = ModeSAT | ModePB | ModeWBO | ModeMaxSAT | ModeMIP
    deriving (Eq, Ord, Show)
  
  data Options = Options
    { optInput :: FilePath
    , optMode :: Maybe Mode
    , optSolver :: String
    , optPrintRational :: Bool
    , optWriteFile :: Maybe FilePath
    , optNoMIP :: Bool
    , optPivotStrategy :: Simplex.PivotStrategy -- String
    , optBoundTightening :: Bool
    , optNThread :: Int
    , optOmegaReal :: String
    , optFileEncoding :: Maybe String
    , optMaxSATCompactVLine :: Bool
    , optPBFastParser :: Bool
    , optExitCode :: Bool
    } deriving (Eq, Show)
  
  optionsParser :: Parser Options
  optionsParser = Options
    <$> fileInput
    <*> modeOption
    <*> solverOption
    <*> printRationalOption
    <*> writeFileOption
    <*> noMIPOption
    <*> pivotStrategyOption
    <*> boundTighteningOption
    <*> nThreadOption
    <*> omegaRealOption
    <*> fileEncodingOption
    <*> maxsatCompactVLineOption
    <*> pbFastParserOption
    <*> exitCodeOption
    where
      fileInput :: Parser FilePath
      fileInput = argument str (metavar "FILE")
  
      modeOption :: Parser (Maybe Mode)
      modeOption = optional $
            flag' ModeSAT    (long "sat"    <> help "solve boolean satisfiability problem in .cnf file")
        <|> flag' ModePB     (long "pb"     <> help "solve pseudo boolean problem in .opb file")
        <|> flag' ModeWBO    (long "wbo"    <> help "solve weighted boolean optimization problem in .wbo file")
        <|> flag' ModeMaxSAT (long "maxsat" <> help "solve MaxSAT problem in .cnf or .wcnf file")
        <|> flag' ModeMIP    (long "mip"    <> help "solve LP/MIP problem in .lp or .mps file")
        <|> flag' ModeMIP    (long "lp"     <> help "alias for --mip (deprecated)")
  
      solverOption :: Parser String
      solverOption = strOption
        $  long "solver"
        <> metavar "SOLVER"
        <> help "Solver algorithm: mip, omega-test, cooper, cad, old-mip, ct"
        <> value "mip"
        <> showDefaultWith id
  
      printRationalOption :: Parser Bool
      printRationalOption = switch
        $  long "print-rational"
        <> help "print rational numbers instead of decimals"
  
      writeFileOption :: Parser (Maybe FilePath)
      writeFileOption = optional $ strOption
        $  short 'w'
        <> metavar "FILE"
        <> help "write solution to filename in Gurobi .sol format"
  
      noMIPOption :: Parser Bool
      noMIPOption = switch
        $  long "nomip"
        <> help "consider all integer variables as continuous"
  
      pivotStrategyOption :: Parser Simplex.PivotStrategy
      pivotStrategyOption = option (maybeReader Simplex.parsePivotStrategy)
        $  long "pivot-strategy"
        <> metavar "NAME"
        <> help ("pivot strategy for simplex: " ++ intercalate ", " [Simplex.showPivotStrategy ps | ps <- [minBound..maxBound]])
        <> value (Simplex.configPivotStrategy def)
        <> showDefaultWith Simplex.showPivotStrategy
  
      boundTighteningOption :: Parser Bool
      boundTighteningOption =  switch
        $  long "bound-tightening"
        <> help "enable bound tightening in simplex algorithm"
  
      nThreadOption :: Parser Int
      nThreadOption = option auto
        $  long "threads"
        <> metavar "INT"
        <> help "number of threads to use (0: auto)"
        <> value 0
        <> showDefault
  
      omegaRealOption :: Parser String
      omegaRealOption = strOption
        $  long "omega-real"
        <> metavar "SOLVER"
        <> help "fourier-motzkin, virtual-substitution (or vs), cad, simplex, none"
        <> value "fourier-motzkin"
        <> showDefaultWith id
  
      fileEncodingOption :: Parser (Maybe String)
      fileEncodingOption = optional $ strOption
        $  long "encoding"
        <> metavar "ENCODING"
        <> help "file encoding for LP/MPS files"
  
      maxsatCompactVLineOption :: Parser Bool
      maxsatCompactVLineOption = switch
        $  long "maxsat-compact-v-line"
        <> help "print Max-SAT solution in the new compact v-line format"
  
      pbFastParserOption :: Parser Bool
      pbFastParserOption = switch
        $  long "pb-fast-parser"
        <> help "use attoparsec-based parser instead of megaparsec-based one for speed"
  
      exitCodeOption :: Parser Bool
      exitCodeOption = switch
        $  long "exit-code"
        <> help "exit with an exit code based on solution status (only for SAT and MaxSAT mode)"
  
  parserInfo :: ParserInfo Options
  parserInfo = info (helper <*> versionOption <*> optionsParser)
    $  fullDesc
    <> header "toysolver - a solver for arithmetic problems"
    where
      versionOption :: Parser (a -> a)
      versionOption = infoOption (V.showVersion version)
        $  hidden
        <> long "version"
        <> help "Show version"
  
  -- ---------------------------------------------------------------------------
  
  data SolutionStatus
    = UNKNOWN
    | UNSUPPORTED
    | SATISFIABLE
    | UNSATISFIABLE
    | OPTIMUM_FOUND
    | UNBOUNDED
    deriving (Show, Eq, Ord, Bounded, Enum)
  
  run
    :: String
    -> Options
    -> MIP.Problem Rational
    -> (SolutionStatus -> IO ())
    -> (Map MIP.Var Rational -> IO ())
    -> IO ()
  run solver opt prob setSolutionStatus printModel = do
    unless (Set.null (MIP.semiContinuousVariables prob)) $ do
      hPutStrLn stderr "semi-continuous variables are not supported."
      exitFailure
    case map toLower solver of
      s | s `elem` ["omega", "omega-test", "cooper"] -> solveByQE
      s | s `elem` ["old-mip"] -> solveByMIP
      s | s `elem` ["cad"] -> solveByCAD
      s | s `elem` ["ct", "conti-traverso"] -> solveByContiTraverso
      _ -> solveByMIP2
    where
      vs = MIP.variables prob
      vsAssoc = zip (Set.toList vs) [0..]
      nameToVar = Map.fromList vsAssoc
      varToName = IntMap.fromList [(v,name) | (name,v) <- vsAssoc]
  
      compileE :: MIP.Expr Rational -> Expr Rational
      compileE = foldr (+) (Const 0) . map compileT . MIP.terms
  
      compileT :: MIP.Term Rational -> Expr Rational
      compileT (MIP.Term c vs3) =
        foldr (*) (Const c) [Var (nameToVar Map.! v) | v <- vs3]
  
      obj = compileE $ MIP.objExpr $ MIP.objectiveFunction prob
  
      cs1 = do
        v <- Set.toList vs
        let v2 = Var (nameToVar Map.! v)
        let (l,u) = MIP.getBounds prob v
        [Const x .<=. v2 | MIP.Finite x <- return l] ++
          [v2 .<=. Const x | MIP.Finite x <- return u]
      cs2 = do
        MIP.Constraint
          { MIP.constrIndicator = ind
          , MIP.constrExpr = e
          , MIP.constrLB = lb
          , MIP.constrUB = ub
          } <- MIP.constraints prob
        case ind of
          Nothing -> do
            let e2 = compileE e
            msum
              [ case lb of
                  MIP.NegInf -> []
                  MIP.PosInf -> [OrdRel 1 Le 0] -- False
                  MIP.Finite x -> [OrdRel e2 Ge (Const x)]
              , case ub of
                  MIP.NegInf -> [OrdRel 1 Le 0] -- False
                  MIP.PosInf -> []
                  MIP.Finite x -> [OrdRel e2 Le (Const x)]
              ]
          Just _ -> error "indicator constraint is not supported yet"
  
      ivs
        | optNoMIP opt = Set.empty
        | otherwise    = MIP.integerVariables prob
  
      vs2  = IntMap.keysSet varToName
      ivs2 = IntSet.fromList . map (nameToVar Map.!) . Set.toList $ ivs
  
      solveByQE =
        case mapM LAFOL.fromFOLAtom (cs1 ++ cs2) of
          Nothing -> do
            setSolutionStatus UNKNOWN
          Just cs ->
            case f vs2 cs ivs2 of
              Nothing -> do
                setSolutionStatus UNSATISFIABLE
              Just m -> do
                putOLine $ showValue (FOL.evalExpr m obj)
                setSolutionStatus SATISFIABLE
                let m2 = Map.fromAscList [(v, m IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
                printModel m2
         where
           f = case solver of
                 "omega"      -> OmegaTest.solveQFLIRAConj omegaOpt
                 "omega-test" -> OmegaTest.solveQFLIRAConj omegaOpt
                 "cooper"     -> Cooper.solveQFLIRAConj
                 _ -> error "unknown solver"
  
           omegaOpt =
             def
             { OmegaTest.optCheckReal = realSolver
             }
             where
               realSolver =
                 case optOmegaReal opt of
                   "fourier-motzkin" -> OmegaTest.checkRealByFM
                   "virtual-substitution" -> OmegaTest.checkRealByVS
                   "vs"              -> OmegaTest.checkRealByVS
                   "cad"             -> OmegaTest.checkRealByCAD
                   "simplex"         -> OmegaTest.checkRealBySimplex
                   "none"            -> OmegaTest.checkRealNoCheck
                   s                 -> error ("unknown solver: " ++ s)
  
      solveByMIP = do
        let prob2 = do
              cs'  <- mapM LAFOL.fromFOLAtom (cs1 ++ cs2)
              obj' <- LAFOL.fromFOLExpr obj
              return (cs',obj')
        case prob2 of
          Nothing -> do
            setSolutionStatus UNSUPPORTED
          Just (cs',obj') ->
            case TextbookMIP.optimize (MIP.objDir $ MIP.objectiveFunction prob) obj' cs' ivs2 of
              TextbookMIP.OptUnsat -> do
                setSolutionStatus UNSATISFIABLE
              TextbookMIP.Unbounded -> do
                setSolutionStatus UNBOUNDED
              TextbookMIP.Optimum r m -> do
                putOLine $ showValue r
                setSolutionStatus OPTIMUM_FOUND
                let m2 = Map.fromAscList [(v, m IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
                printModel m2
  
      solveByMIP2 = do
        simplex <- Simplex.newSolver
  
        let config =
              def
              { Simplex.configPivotStrategy = optPivotStrategy opt
              , Simplex.configEnableBoundTightening = optBoundTightening opt
              }
            nthreads = optNThread opt
  
        Simplex.setConfig simplex config
        Simplex.setLogger simplex putCommentLine
        Simplex.enableTimeRecording simplex
        replicateM_ (length vsAssoc) (Simplex.newVar simplex) -- XXX
        Simplex.setOptDir simplex $ MIP.objDir $ MIP.objectiveFunction prob
        Simplex.setObj simplex $ fromJust (LAFOL.fromFOLExpr obj)
        putCommentLine "Loading constraints... "
        forM_ (cs1 ++ cs2) $ \c ->
          Simplex.assertAtom simplex $ fromJust (LAFOL.fromFOLAtom c)
        putCommentLine "Loading constraints finished"
  
        mip <- MIPSolver.newSolver simplex ivs2
        MIPSolver.setShowRational mip printRat
        MIPSolver.setLogger mip putCommentLine
        MIPSolver.setOnUpdateBestSolution mip $ \_m val -> putOLine (showValue val)
  
        procs <-
          if nthreads >= 1
          then return nthreads
          else do
            ncap  <- getNumCapabilities
            procs <- getNumProcessors
            return $ max (procs - 1) ncap
        setNumCapabilities procs
        MIPSolver.setNThread mip procs
  
        ret <- MIPSolver.optimize mip
        case ret of
          Simplex.Unsat -> do
            setSolutionStatus UNSATISFIABLE
          Simplex.Unbounded -> do
            setSolutionStatus UNBOUNDED
            Just m <- MIPSolver.getBestModel mip
            let m2 = Map.fromAscList [(v, m IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
            printModel m2
          Simplex.Optimum -> do
            Just m <- MIPSolver.getBestModel mip
            setSolutionStatus OPTIMUM_FOUND
            let m2 = Map.fromAscList [(v, m IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
            printModel m2
          Simplex.ObjLimit -> do
            error "should not happen"
  
      solveByCAD
        | not (IntSet.null ivs2) = do
            setSolutionStatus UNSUPPORTED
            putCommentLine "integer variables are not supported by CAD"
        | otherwise = do
            let cs = map (fmap f) $ cs1 ++ cs2
                vs3 = Set.fromAscList $ IntSet.toAscList vs2
            case CAD.solve vs3 cs of
              Nothing -> do
                setSolutionStatus UNSATISFIABLE
              Just m -> do
                let m2 = IntMap.map (\x -> AReal.approx x (2^^(-64::Int))) $
                           IntMap.fromAscList $ Map.toAscList $ m
                putOLine $ showValue (FOL.evalExpr m2 obj)
                setSolutionStatus SATISFIABLE
                let m3 = Map.fromAscList [(v, m2 IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
                printModel m3
        where
          f (Const r)   = P.constant r
          f (Var v)     = P.var v
          f (e1 :+: e2) = f e1 + f e2
          f (e1 :*: e2) = f e1 * f e2
          f (e1 :/: e2)
            | P.deg p > 0 = error "can't handle rational expression"
            | otherwise   = P.mapCoeff (/ c) $ f e1
            where
              p = f e2
              c = P.coeff P.mone p
  
      solveByContiTraverso
        | not (vs `Set.isSubsetOf` ivs) = do
            setSolutionStatus UNSUPPORTED
            putCommentLine "continuous variables are not supported by Conti-Traverso algorithm"
        | otherwise = do
            let tmp = do
                  linObj <- LAFOL.fromFOLExpr obj
                  linCon <- mapM LAFOL.fromFOLAtom (cs1 ++ cs2)
                  return (linObj, linCon)
            case tmp of
              Nothing -> do
                setSolutionStatus UNSUPPORTED
                putCommentLine "non-linear expressions are not supported by Conti-Traverso algorithm"
              Just (linObj, linCon) ->
                case ContiTraverso.solve P.grlex vs2 (MIP.objDir $ MIP.objectiveFunction prob) linObj linCon of
                  Nothing -> do
                    setSolutionStatus UNSATISFIABLE
                  Just m -> do
                    let m2 = IntMap.map fromInteger m
                    putOLine $ showValue (FOL.evalExpr m2 obj)
                    setSolutionStatus OPTIMUM_FOUND
                    let m3 = Map.fromAscList [(v, m2 IntMap.! (nameToVar Map.! v)) | v <- Set.toList vs]
                    printModel m3
  
      printRat :: Bool
      printRat = optPrintRational opt
  
      showValue :: Rational -> String
      showValue = showRational printRat
  
  mipPrintModel :: Handle -> Bool -> Map MIP.Var Rational -> IO ()
  mipPrintModel h asRat m =
    forM_ (Map.toList m) $ \(v, val) ->
      hPrintf h "v %s = %s\n" (MIP.varName v) (showRational asRat val)
  
  putCommentLine :: String -> IO ()
  putCommentLine s = do
    putStr "c "
    putStrLn s
    hFlush stdout
  
  putSLine :: String -> IO ()
  putSLine  s = do
    putStr "s "
    putStrLn s
    hFlush stdout
  
  putOLine :: String -> IO ()
  putOLine  s = do
    putStr "o "
    putStrLn s
    hFlush stdout
  
  -- ---------------------------------------------------------------------------
  
  main :: IO ()
  main = do
  #ifdef FORCE_CHAR8
    setEncodingChar8
  #endif
  
    opt <- execParser parserInfo
  
    ref <- newIORef Nothing
    let putStatus status = do
          m <- readIORef ref
          case m of
            Just _ -> error "solution status is already set"
            Nothing -> writeIORef ref (Just status)
  
        setSolutionStatusSAT status = do
          putStatus status
          case status of
            UNKNOWN       -> putSLine "UNKNOWN"
            UNSUPPORTED   -> putSLine "UNKNOWN"
            SATISFIABLE   -> putSLine "SATISFIABLE"  
            UNSATISFIABLE -> putSLine "UNSATISFIABLE"
            OPTIMUM_FOUND -> putSLine "SATISFIABLE"
            UNBOUNDED     -> putSLine "UNKNOWN"
  
        setSolutionStatusMaxSAT status = do
          putStatus status
          case status of
            UNKNOWN       -> putSLine "UNKNOWN"
            UNSUPPORTED   -> putSLine "UNKNOWN"
            SATISFIABLE   -> putSLine "SATISFIABLE"  
            UNSATISFIABLE -> putSLine "UNSATISFIABLE"
            OPTIMUM_FOUND -> putSLine "OPTIMUM FOUND"
            UNBOUNDED     -> putSLine "UNKNOWN"
  
        setSolutionStatusPB status = do
          putStatus status
          case status of
            UNKNOWN       -> putSLine "UNKNOWN"
            UNSUPPORTED   -> putSLine "UNSUPPORTED"
            SATISFIABLE   -> putSLine "SATISFIABLE"  
            UNSATISFIABLE -> putSLine "UNSATISFIABLE"
            OPTIMUM_FOUND -> putSLine "OPTIMUM FOUND"
            UNBOUNDED     -> putSLine "UNKNOWN"
  
    case fromMaybe ModeMIP (optMode opt) of
      ModeSAT -> do
        cnf <- FF.readFile (optInput opt)
        let (mip,info2) = sat2ip cnf
  
        run (optSolver opt) opt (fmap fromInteger mip) setSolutionStatusSAT $ \m -> do
          let m2 = transformBackward info2 m
          satPrintModel stdout m2 0
          writeSOLFileSAT opt m2
  
        when (optExitCode opt) $ do
          status <- readIORef ref
          case fromMaybe UNKNOWN status of
            UNKNOWN       -> exitSuccess
            UNSUPPORTED   -> exitSuccess
            SATISFIABLE   -> exitWith (ExitFailure 10)
            UNSATISFIABLE -> exitWith (ExitFailure 20)
            OPTIMUM_FOUND -> exitWith (ExitFailure 10)
            UNBOUNDED     -> exitSuccess
  
      ModePB -> do
        pb <-
          if optPBFastParser opt then
            liftM FF.unWithFastParser $ FF.readFile (optInput opt)
          else
            FF.readFile (optInput opt)
        let (mip,info2) = pb2ip pb
        run (optSolver opt) opt (fmap fromInteger mip) setSolutionStatusPB $ \m -> do
          let m2 = transformBackward info2 m
          pbPrintModel stdout m2 0
          writeSOLFileSAT opt m2
  
      ModeWBO -> do
        wbo <-
          if optPBFastParser opt then
            liftM FF.unWithFastParser $ FF.readFile (optInput opt)
          else
            FF.readFile (optInput opt)
        let (mip,info2) = wbo2ip False wbo
        run (optSolver opt) opt (fmap fromInteger mip) setSolutionStatusPB $ \m -> do
          let m2 = transformBackward info2 m
          pbPrintModel stdout m2 0
          writeSOLFileSAT opt m2
  
      ModeMaxSAT -> do
        wcnf <- FF.readFile (optInput opt)
        let (mip,info2) = maxsat2ip False wcnf
        run (optSolver opt) opt (fmap fromInteger mip) setSolutionStatusMaxSAT $ \m -> do
          let m2 = transformBackward info2 m
          if optMaxSATCompactVLine opt then
            maxsatPrintModelCompact stdout m2 0
          else
            maxsatPrintModel stdout m2 0
          writeSOLFileSAT opt m2
  
        when (optExitCode opt) $ do
          status <- readIORef ref
          case fromMaybe UNKNOWN status of
            UNKNOWN       -> exitSuccess
            UNSUPPORTED   -> exitSuccess
            SATISFIABLE   -> exitWith (ExitFailure 10)
            UNSATISFIABLE -> exitWith (ExitFailure 20)
            OPTIMUM_FOUND -> exitWith (ExitFailure 30)
            UNBOUNDED     -> exitSuccess
  
      ModeMIP -> do
        enc <- T.mapM mkTextEncoding $ optFileEncoding opt
        mip <- MIP.readFile def{ MIP.optFileEncoding = enc } (optInput opt)
        run (optSolver opt) opt (fmap toRational mip) setSolutionStatusPB $ \m -> do
          mipPrintModel stdout (optPrintRational opt) m
          writeSOLFileMIP opt m
  
    status <- readIORef ref
    case fromMaybe UNKNOWN status of
      UNKNOWN       -> exitFailure
      UNSUPPORTED   -> exitFailure
      _ -> exitSuccess
  
  -- FIXME: 目的関数値を表示するように
  writeSOLFileMIP :: Options -> Map MIP.Var Rational -> IO ()
  writeSOLFileMIP opt m = do
    let sol = MIP.Solution
              { MIP.solStatus = MIP.StatusUnknown
              , MIP.solObjectiveValue = Nothing
              , MIP.solVariables = Map.fromList [(v, Scientific.fromFloatDigits (fromRational val :: Double)) | (v,val) <- Map.toList m]
              }
    writeSOLFileRaw opt sol
  
  -- FIXME: 目的関数値を表示するように
  writeSOLFileSAT :: Options -> SAT.Model -> IO ()
  writeSOLFileSAT opt m = do
    let sol = MIP.Solution
              { MIP.solStatus = MIP.StatusUnknown
              , MIP.solObjectiveValue = Nothing
              , MIP.solVariables = Map.fromList [(fromString ("x" ++ show x), if b then 1 else 0) | (x,b) <- assocs m]
              }
    writeSOLFileRaw opt sol
  
  writeSOLFileRaw :: Options -> MIP.Solution Scientific -> IO ()
  writeSOLFileRaw opt sol =
    case optWriteFile opt of
      Just fname -> GurobiSol.writeFile fname sol
      Nothing -> return ()
  
  