{-# LANGUAGE CPP #-}

#ifndef CYCLES
# define CYCLES 2
#endif

#ifndef ROUNDS
# define ROUNDS 2
#endif

-- Set up satisfiability problem (in smtlib format, LRA) for a random
-- sequence of pivoting steps that lead to a potential cycle.
--
-- We restrict ourselves to tableaux with 2 rows, so pivoting steps
-- must alternate between the first and the secon row. Using the full
-- tableaux, every pivoting step is a row operation, which can be
-- represented by a matrices [a,0; b,1] and [1,a;0,b], alternating
-- between these two shapes, subject to the constraint that there
-- is an even number of steps and the product of these matrices is
-- the identity matrix.
--
-- To this end, we set up a sequence of random invertible matrices of
-- suitable shape for the first n rounds, and fill it up with 2 or 4
-- more matrices to obtain the identity matrix for the product.
--
-- The tableaux can be computed from the sequence of steps by noting
-- that it has n+2 (or n+4) columns, and that after each pivoting step,
-- the corresponding column must be a unit vector, alternating between
-- [1;0] and [0;1].
--
-- In order to find a suitable initial assignment of the variables, the
-- generated SMT problem encodes the constraints for pivot selection,
-- namely that the constraint of the row being pivoted is not satisfied,
-- and that the selected coefficient has the right sign.
--
-- Note that no effort is made to keep solutions small... it appears
-- that yices does a reasonable job of that by itself.
--
-- Basic usage:
--   ./simplex > test.smt; yices-smt2 test.smt
--
-- author: Bertram Felgenhauer
-- code:     Apr 2016
-- comments: Nov 2016

import Data.Ratio

import System.Random
import Control.Monad
import Control.Applicative

-- number of cycles
cycles :: Int
cycles = CYCLES

-- number of rounds is cycles*(rounds + 2) (sometimes + 0 or + 4)
rounds :: Int
rounds = ROUNDS

------------------------------------------------------------------------
-- fundamentals

-- scalars
type N = Rational

-- column vector
data V2 = V2 !N !N
    deriving (Eq, Show)

-- unit vectors
e1, e2 :: V2
e1 = V2 1 0
e2 = V2 0 1

-- scalar product
(|*|) :: V2 -> V2 -> N
V2 a b |*| V2 c d = a*c + b*d

-- additive group
instance Num V2 where
    V2 a b + V2 c d = V2 (a + c) (b + d)
    V2 a b - V2 c d = V2 (a - c) (b - d)
    _ * _ = error "V2::(*)"
    negate (V2 a b) = V2 (negate a) (negate b)
    fromInteger = error "V2::fromInteger"
    abs = error "V2::abs"
    signum = error "V2::signum"

-- square matrix
data M2 = M2 !N !N !N !N
    deriving (Eq, Show)

-- multiply matrix by vector
(*|) :: M2 -> V2 -> V2
M2 a b c d *| V2 e f = V2 (a*e + b*f) (c*e + d*f)

-- make matrix from column vectors
m2 :: V2 -> V2 -> M2
m2 (V2 a c) (V2 b d) = M2 a b c d

-- matrix semiring
instance Num M2 where
    M2 a b c d + M2 e f g h = M2 (a + e) (b + f) (c + g) (d + h)
    M2 a b c d - M2 e f g h = M2 (a - e) (b - f) (c - g) (d - h)
    m * M2 a b c d = m2 (m *| V2 a c) (m *| V2 b d)
    negate (M2 a b c d) = M2 (negate a) (negate b) (negate c) (negate d)
    fromInteger z = M2 (fromInteger z) 0 0 (fromInteger z)
    abs = error "M2::abs"
    signum = error "M2::signum"

-- invert matrix
inv :: M2 -> Maybe M2
inv (M2 a b c d)
    | det == 0  = Nothing
    | otherwise = Just $ M2 (d*f) (-b*f) (-c*f) (a*f)
  where
    det = a*d - b*c
    f = 1/det

------------------------------------------------------------------------
-- tableaux generation

-- sequence of matrices alternating between [1,a;0,b] and [a,0;b,1]
mkMs :: [V2] -> [M2]
mkMs xs | odd (length xs) = error "only even number of rounds is supported"
mkMs xs = zipWith ($) (cycle [m2 e1, flip m2 e2]) xs

-- check for invertibility of mkMs xs
mkVs :: [V2] -> [V2]
mkVs xs = case inv (product (mkMs xs)) of
    Just mi -> xs ++ factor mi
    _ -> error "not invertible"

-- factor m@M2 a b c d into row operations (typically 2, but sometimes 0 or 4)
-- * return [u,v] with P<u,v> = m2 e1 u * m2 v e2 = m
-- * exceptional case: d=0, c/=0: return [u,v,w,x] with P<u,v> * P<w,x> = m
factor :: M2 -> [V2]
factor (M2 1 0 0 1) = []
factor (M2 a b 0 0) = [V2 b 0, V2 a 0]
factor (M2 a b c 0) = [V2 (a+b) c, V2 (-b) 1, V2 (-1) 1, V2 1 0]
factor (M2 a b c d) = [V2 b d, V2 (a - b*c/d) (c/d)]

-- compute tableaux from sequence of vectors
mkM :: [V2] -> ([V2],[Bool])
mkM vs
    | product ms == m2 e1 e2 = foldr step (const ([], [])) ms (negate e2)
    | otherwise = error "mkM: product must be the unit matrix"
  where
    ms :: [M2]
    ms = mkMs vs
    step :: M2 -> (V2 -> ([V2],[Bool])) -> V2 -> ([V2],[Bool])
    step m r v = let (vs, ds) = r (m2 e2 e1 *| v) in
         (map (m *|) vs ++ [v], ds ++ [(m *| v) |*| v < 0])

------------------------------------------------------------------------------
-- ad-hoc SMT pretty printing for SMT-lib encoding

type Doc = String

-- s-expression from list
slist :: [Doc] -> Doc
slist = ("(" ++) . (++ ")") . unwords

-- s-expression in apply from
s :: Doc -> [Doc] -> Doc
s o = slist . (o :)

-- basic type
stype :: Doc
stype = "Real" -- using LRA

-- format integer
sint :: Integer -> Doc
sint x
    | x < 0 = s "-" [sint (-x)]
    | otherwise = show x

-- format rational number
srat :: Rational -> Doc
srat x
    | x < 0 = s "-" [srat (-x)]
    | denominator x == 1 = sint (numerator x)
    | otherwise = s "/" [sint (numerator x), sint (denominator x)]

-- format scalar
sn :: N -> Doc
sn = srat

-- declare variable
sdecl :: String -> Doc
sdecl v = s "declare-fun" [v, "()", stype]

-- declare an assignment; at any point in time, the two equations given by
-- the original tableaux must be satisfied, so we add those as constraints
declVars :: [V2] -> String -> [Doc]
declVars m suf =
    [sdecl (v i) | (_, i) <- zs] ++
    [s "assert" [s "=" ["0", s "+" [s "*" [sn c, v i] | (V2 c _, i) <- zs]]]] ++
    [s "assert" [s "=" ["0", s "+" [s "*" [sn c, v i] | (V2 _ c, i) <- zs]]]]
  where
    zs = zip m [0..]
    v i = "x" ++ show i ++ suf

-- furthermore, in each round, the variables not being pivoted remain
-- unchanged, and the pivoted variable moves to its upper or lower bound,
-- depending on which bound was violated before the step; we also
-- assert that one of the bounds must be violated
encodeRound :: [V2] -> [Bool] -> Int -> [Doc]
encodeRound m ds rd =
    declVars m ('r' : show (rd+1)) ++
    [s "assert" [s "and" [s "=" [vx i 1, vx i 0] | i <- [1..l-3]]]] ++
    [s "assert" [s "ite" [c True, a True,
                 s "ite" [c False, a False, "false"]]]] ++
--    [s "echo" [show $ "- round " ++ show (rd + 1)],
--     s "check-sat" []] ++
    []
  where
    old = 0
    new = -2
    c f = if f then s "<" [vx new 0, vl new] else s "<" [vu new, vx new 0]
    a f = s "and" [
        if f == ds !! (rd `mod` l) then s "<" [vx old 0, vu old]
        else s "<" [vl old, vx old 0],
        if f then s "=" [vx new 1, vl new] else s "=" [vx new 1, vu new]]
    l = length ds
    vl i = 'l':show ((rd + i) `mod` l)
    vu i = 'u':show ((rd + i) `mod` l)
    vx i j = "x" ++ show ((rd + i) `mod` l) ++ "r" ++ show (rd + j)


-- full encoding
encode :: [V2] -> String
encode vs = unlines $
    [-- preamble
     s "set-option" [":produce-models", "true"],
     s "set-option" [":print-success", "false"],
     s "set-logic" ["QF_LRA"]] ++
     -- declare variables for lower and upper bounds
     map unwords [
         [sdecl ('l':i), sdecl ('u':i), s "assert" [s "<=" ['l':i, 'u':i]]]
         | i <- map show [0..l-1]
     ] ++
     -- declare initial assignment
     declVars m "r0" ++
     -- declare assignments for each round
     [s "assert" [s "and" [s "<=" [vl i, vx i 0], s "<=" [vx i 0, vu i]]]
     | i <- [0..l-3]] ++
     concat [encodeRound m ds i | i <- [0..rds-1]] ++
     -- assert that the final assignment equals the initial one
     [s "assert" [s "and" [s "=" [vx i 0, vx i rds],
                           s "=" [vx i 0, sn 0]]] | i <- [0..l-1]] ++
     -- ensure satisfiability of the tableaux; this is optional
     declVars m "" ++
     [s "assert" [s "and" [s "<=" [vl i, vx' i],
                           s "<=" [vx' i, vu i]]] | i <- [0..l-1]] ++
     -- check satisfiability and print assignments
     [s "check-sat" [],
      s "get-value" [slist (map vl [0..l-1])],
      s "get-value" [slist (map vu [0..l-1])]] ++
     concat [
      [s "get-value" [slist [vx i j | i <- [0..l-1]]]]
      | j <- [0..out]] ++
     [s "get-value" [slist [vx' i | i <- [0..l-1]]]] ++
     []
  where
    rds = cycles * l
    out = rds
    l = length vs
    (m, ds) = mkM vs
    vl i = 'l':show i
    vu i = 'u':show i
    vx i j = "x" ++ show i ++ "r" ++ show j
    vx' i = "x" ++ show i

-- produce encoding: compute random matrixes
main = do
    vs <- replicateM rounds rv
    mapM_ (putStrLn . ("; " ++) . showM) $ mkMs $ mkVs vs
    putStrLn $ "; [" ++ showM' (fst $ mkM $ mkVs vs) ++ "] " ++ show (snd $ mkM $ mkVs vs)
    putStr $ encode $ mkVs vs
  where
    showM (M2 a b c d) = concat ["[", sn a, ",",  sn b, ";", sn c, ",", sn d, "]"]
    showM' [] = "";
    showM' [V2 a b] = sn a ++ "," ++ sn b
    showM' (v:vs) = showM' [v] ++ ";" ++ showM' vs
    sn x | den == 1 = show num | otherwise = show num ++ "/" ++ show den where
        (num, den) = (numerator x, denominator x)
    -- matrix coefficients are selected as rational numbers with
    -- numerator and denominator up to 3
    cs = [a % b | a <- [-3..3], a /= 0, b <- [1..3], gcd a b == 1]
    rc = (cs !!) <$> randomRIO (0, length cs - 1)
    rv = V2 <$> rc <*> rc

-- test data (unused)
vs0 :: [V2]
vs0 = [V2 (-9) (-3), V2 (1/50) (1/150), V2 10500 (-40),
       V2 (-25/8) (1/160), V2 (-8) (1/30), V2 (-4) (-2)]
    -- survives 17 rounds