{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Text.SDPFile
-- Copyright   :  (c) Masahiro Sakai 2012,2016
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- References:
--
-- * SDPA (Semidefinite Programming Algorithm) User's Manual
--   <http://sdpa.indsys.chuo-u.ac.jp/~fujisawa/sdpa_doc.pdf>
--
-- * <http://euler.nmt.edu/~brian/sdplib/FORMAT>
--
-----------------------------------------------------------------------------
module ToySolver.Text.SDPFile
  ( -- * The problem type
    Problem (..)
  , Matrix
  , Block
  , mDim
  , nBlock
  , blockElem
    -- * The solution type
  , Solution (..)
  , evalPrimalObjective
  , evalDualObjective

    -- * File I/O
  , readDataFile
  , writeDataFile

    -- * Construction
  , DenseMatrix
  , DenseBlock
  , denseMatrix
  , denseBlock
  , diagBlock

    -- * Rendering
  , renderData
  , renderSparseData

    -- * Parsing
  , ParseError
  , parseData
  , parseSparseData
  ) where

import Control.Exception
import Control.Monad
import qualified Data.ByteString.Lazy as BL
import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Builder.Scientific as B
import Data.Char
import qualified Data.Foldable as F
import Data.List (intersperse)
import Data.Scientific (Scientific)
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import Data.Void
import Data.Word
import System.FilePath (takeExtension)
import System.IO
import qualified Text.Megaparsec as MegaParsec
import Text.Megaparsec hiding (ParseError, oneOf)
import Text.Megaparsec.Byte
import qualified Text.Megaparsec.Byte.Lexer as Lexer

type C e s m = (MonadParsec e s m, Token s ~ Word8)
type ParseError = MegaParsec.ParseErrorBundle BL.ByteString Void

anyChar :: C e s m => m Word8
anyChar = anySingle

-- ---------------------------------------------------------------------------
-- problem description
-- ---------------------------------------------------------------------------

data Problem
  = Problem
  { blockStruct :: [Int]      -- ^ the block strcuture vector (bLOCKsTRUCT)
  , costs       :: [Scientific] -- ^ Constant Vector
  , matrices    :: [Matrix]   -- ^ Constraint Matrices
  }
  deriving (Show, Ord, Eq)

type Matrix = [Block]

type Block = Map (Int,Int) Scientific

-- | the number of primal variables (mDim)
mDim :: Problem -> Int
mDim prob = length (matrices prob) - 1

-- | the number of blocks (nBLOCK)
nBlock :: Problem -> Int
nBlock prob = length (blockStruct prob)

blockElem :: Int -> Int -> Block -> Scientific
blockElem i j b = Map.findWithDefault 0 (i,j) b

-- ---------------------------------------------------------------------------
-- solution
-- ---------------------------------------------------------------------------

data Solution
  = Solution
  { primalVector :: [Scientific] -- ^ The primal variable vector x
  , primalMatrix :: Matrix -- ^ The primal variable matrix X
  , dualMatrix   :: Matrix -- ^ The dual variable matrix Y
  }
  deriving (Show, Ord, Eq)

evalPrimalObjective :: Problem -> Solution -> Scientific
evalPrimalObjective prob sol = sum $ zipWith (*) (costs prob) (primalVector sol)

evalDualObjective :: Problem -> Solution -> Scientific
evalDualObjective Problem{ matrices = [] } _ = error "evalDualObjective: invalid problem data"
evalDualObjective Problem{ matrices = f0:_ } sol =
  sum $ zipWith (\blk1 blk2 -> F.sum (Map.intersectionWith (*) blk1 blk2)) f0 (dualMatrix sol)

-- ---------------------------------------------------------------------------
-- construction
-- ---------------------------------------------------------------------------

type DenseMatrix = [DenseBlock]

type DenseBlock = [[Scientific]]

denseBlock :: DenseBlock -> Block
denseBlock xxs = Map.fromList [((i,j),x) | (i,xs) <- zip [1..] xxs, (j,x) <- zip [1..] xs, x /= 0]

denseMatrix :: DenseMatrix -> Matrix
denseMatrix = map denseBlock

diagBlock :: [Scientific] -> Block
diagBlock xs = Map.fromList [((i,i),x) | (i,x) <- zip [1..] xs]

-- ---------------------------------------------------------------------------
-- File I/O
-- ---------------------------------------------------------------------------

-- | Parse a SDPA format file (.dat) or a SDPA sparse format file (.dat-s)..
readDataFile :: FilePath -> IO Problem
readDataFile fname = do
  p <- case map toLower (takeExtension fname) of
    ".dat" -> return pDataFile
    ".dat-s" -> return pSparseDataFile
    ext -> ioError $ userError $ "unknown extension: " ++ ext
  s <- BL.readFile fname
  case parse (p <* eof) fname s of
    Left e -> throw (e :: ParseError)
    Right prob -> return prob

writeDataFile :: FilePath -> Problem -> IO ()
writeDataFile fname prob = do
  isSparse <- case map toLower (takeExtension fname) of
    ".dat" -> return False
    ".dat-s" -> return True
    ext -> ioError $ userError $ "unknown extension: " ++ ext
  withBinaryFile fname WriteMode $ \h -> do
    B.hPutBuilder h $ renderImpl isSparse prob

-- ---------------------------------------------------------------------------
-- parsing
-- ---------------------------------------------------------------------------

-- | Parse a SDPA format (.dat) string.
parseData :: String -> BL.ByteString -> Either ParseError Problem
parseData = parse (pDataFile <* eof)

-- | Parse a SDPA sparse format (.dat-s) string.
parseSparseData :: String -> BL.ByteString -> Either ParseError Problem
parseSparseData = parse (pSparseDataFile <* eof)

pDataFile :: C e s m => m Problem
pDataFile = do
  _ <- many pComment
  m  <- nat_line -- mDim
  _n <- nat_line -- nBlock
  bs <- pBlockStruct -- bLOCKsTRUCT
  cs <- pCosts
  ms <- pDenseMatrices (fromIntegral m) bs
  space
  return $
    Problem
    { blockStruct = bs
    , costs       = cs
    , matrices    = ms
    }

pSparseDataFile :: C e s m => m Problem
pSparseDataFile = do
  _ <- many pComment
  m  <- nat_line -- mDim
  _n <- nat_line -- nBlock
  bs <- pBlockStruct -- bLOCKsTRUCT
  cs <- pCosts
  ms <- pSparseMatrices (fromIntegral m) bs
  space
  return $
    Problem
    { blockStruct = bs
    , costs       = cs
    , matrices    = ms
    }

pComment :: C e s m => m BL.ByteString
pComment = do
  c <- oneOf "*\""
  cs <- manyTill anyChar newline
  return $ BL.pack (c:cs)

pBlockStruct :: C e s m => m [Int]
pBlockStruct = do
  _ <- optional sep
  let int' = int >>= \i -> optional sep >> return i
  xs <- many int'
  _ <- manyTill anyChar newline
  return $ map fromIntegral xs
  where
    sep = some (oneOf " \t(){},")

pCosts :: C e s m => m [Scientific]
pCosts = do
  let sep = some (oneOf " \t(){},")
      real' = real >>= \r -> optional sep >> return r
  _ <- optional sep
  cs <- many real'
  _ <- newline
  return cs

pDenseMatrices :: C e s m => Int -> [Int] -> m [Matrix]
pDenseMatrices m bs = optional sep >> replicateM (fromIntegral m + 1) pDenceMatrix
  where
    sep = some ((spaceChar >> return ()) <|> (oneOf "(){}," >> return ()))
    real' = real >>= \r -> optional sep >> return r
    pDenceMatrix = forM bs $ \b ->
      if b >= 0
      then do
        xs <- replicateM b (replicateM b real')
        return $ denseBlock xs
      else do
        xs <- replicateM (abs b) real'
        return $ diagBlock xs

pSparseMatrices :: C e s m => Int -> [Int] -> m [Matrix]
pSparseMatrices m bs = do
  xs <- many pLine
  let t = IntMap.unionsWith (IntMap.unionWith Map.union)
            [ IntMap.singleton matno (IntMap.singleton blkno (Map.fromList [((i,j),e),((j,i),e)]))
            | (matno,blkno,i,j,e) <- xs ]
  return $
    [ [IntMap.findWithDefault Map.empty blkno mat | blkno <- [1 .. length bs]]
    | matno <- [0..m], let mat = IntMap.findWithDefault IntMap.empty matno t
    ]

  where
    sep = some (oneOf " \t") >> return ()
    pLine = do
      _ <- optional sep
      matno <- nat
      sep
      blkno <- nat
      sep
      i <- nat
      sep
      j <- nat
      sep
      e <- real
      _ <- optional sep
      _ <- newline
      return (fromIntegral matno, fromIntegral blkno, fromIntegral i, fromIntegral j, e)

nat_line :: C e s m => m Integer
nat_line = do
  space
  n <- nat
  _ <- manyTill anyChar newline
  return n

nat :: C e s m => m Integer
nat = Lexer.decimal

int :: C e s m => m Integer
int = Lexer.signed (return ()) Lexer.decimal

real :: forall e s m. C e s m => m Scientific
real = Lexer.signed (return ()) Lexer.scientific

oneOf :: C e s m => [Char] -> m Word8
oneOf = MegaParsec.oneOf . map (fromIntegral . fromEnum)

-- ---------------------------------------------------------------------------
-- rendering
-- ---------------------------------------------------------------------------

renderData :: Problem -> Builder
renderData = renderImpl False

renderSparseData :: Problem -> Builder
renderSparseData = renderImpl True

renderImpl :: Bool -> Problem -> Builder
renderImpl sparse prob = mconcat
  [
  -- mDim
    B.intDec (mDim prob) <> " = mDIM\n"

  -- nBlock
  , B.intDec (nBlock prob) <> " = nBlock\n"

  -- blockStruct
  , B.char7 '('
  , sepByS [B.intDec i | i <- blockStruct prob] ", "
  , B.char7 ')'
  , " = bLOCKsTRUCT\n"

  -- costs
  , B.char7 '('
  , sepByS [B.scientificBuilder c | c <- costs prob] ", "
  , ")\n"

  -- matrices
  , if sparse
    then mconcat [renderSparseMatrix matno m | (matno, m) <- zip [0..] (matrices prob)]
    else mconcat $ map renderDenseMatrix (matrices prob)
  ]

  where
    renderSparseMatrix :: Int -> Matrix -> Builder
    renderSparseMatrix matno m =
      mconcat [ B.intDec matno <> B.char7 ' ' <>
                B.intDec blkno <> B.char7 ' ' <>
                B.intDec i <> B.char7 ' ' <>
                B.intDec j <> B.char7 ' ' <>
                B.scientificBuilder e <> B.char7 '\n'
              | (blkno, blk) <- zip [(1::Int)..] m, ((i,j),e) <- Map.toList blk, i <= j ]

    renderDenseMatrix :: Matrix -> Builder
    renderDenseMatrix m =
      "{\n" <>
      mconcat [renderDenseBlock b s <> "\n" | (b,s) <- zip m (blockStruct prob)] <>
      "}\n"

    renderDenseBlock :: Block -> Int -> Builder
    renderDenseBlock b s
      | s < 0 =
          "  " <> renderVec [blockElem i i b | i <- [1 .. abs s]]
      | otherwise =
          "  { " <>
          sepByS [renderRow i | i <- [1..s]] ", " <>
          " }"
      where
        renderRow i = renderVec [blockElem i j b | j <- [1..s]]

renderVec :: [Scientific] -> Builder
renderVec xs =
  B.char7 '{' <>
  sepByS (map B.scientificBuilder xs) ", " <>
  B.char7 '}'

sepByS :: [Builder] -> Builder -> Builder
sepByS xs sep = mconcat $ intersperse sep xs