    {-# LANGUAGE GeneralizedNewtypeDeriving #-}
    module Agda.Utils.Warshall where

    {- construct a graph from constraints

       x + n <= y   becomes   x ---(-n)---> y
       x <= n + y   becomes   x ---(+n)---> y

    the default edge (= no edge is) labelled with infinity

    building the graph involves keeping track of the node names.
    We do this in a finite map, assigning consecutive numbers to nodes.
    -}

    import Control.Applicative
    import Control.Monad.State
    import Data.Maybe -- fromJust
    import Data.Array
    import Data.List
    import Data.Map (Map)
    import qualified Data.Map as Map
    import Test.QuickCheck
    import Agda.Utils.TestHelpers
    import Agda.Syntax.Common
    import Agda.Utils.QuickCheck
    import Agda.Utils.SemiRing

    import Debug.Trace

    type Matrix a = Array (Int,Int) a

    -- assuming a square matrix
    warshall :: SemiRing a => Matrix a -> Matrix a
    warshall a0 = loop r a0 where
      b@((r,c),(r',c')) = bounds a0 -- assuming r == c and r' == c'
      loop k a | k <= r' =
        loop (k+1) (array b [ ((i,j),
                               (a!(i,j)) `oplus` ((a!(i,k)) `otimes` (a!(k,j))))
                            | i <- [r..r'], j <- [c..c'] ])
               | otherwise = a

    -- Warshall's algorithm on a graph represented as an adjacency list.
    type AdjList node edge = Map node [(node, edge)]

    warshallG :: (SemiRing edge, Ord node) => AdjList node edge -> AdjList node edge
    warshallG g = fromMatrix $ warshall m
      where
        nodes = zip (nub $ Map.keys g ++ map fst (concat $ Map.elems g))
                    [0..]
        len   = length nodes
        b     = ((0,0), (len - 1,len - 1))

        edge i j = do
          es <- Map.lookup i g
          foldr oplus Nothing [ Just v | (j', v) <- es, j == j' ]

        m = array b [ ((n, m), edge i j) | (i, n) <- nodes, (j, m) <- nodes ]

        fromMatrix matrix = Map.fromList $ do
          (i, n) <- nodes
          let es = [ (fst (nodes !! m), e)
                   | m <- [0..len - 1]
                   , Just e <- [matrix ! (n, m)]
                   ]
          return (i, es)

    -- edge weight in the graph, forming a semi ring

    data Weight = Finite Int | Infinite
                  deriving (Eq)

    inc :: Weight -> Int -> Weight
    inc Infinite   n = Infinite
    inc (Finite k) n = Finite (k + n)

    instance Show Weight where
      show (Finite i) = show i
      show Infinite   = "."

    instance Ord Weight where
      a <= Infinite = True
      Infinite <= b = False
      Finite a <= Finite b = a <= b

    instance SemiRing Weight where
      oplus = min

      otimes Infinite _ = Infinite
      otimes _ Infinite = Infinite
      otimes (Finite a) (Finite b) = Finite (a + b)

    -- constraints ---------------------------------------------------

    -- nodes of the graph are either
    -- * flexible variables (with identifiers drawn from Int),
    -- * rigid variables (also identified by Ints), or
    -- * constants (like 0, infinity, or anything between)

    data Node = Rigid Rigid
              | Flex  FlexId
                deriving (Eq, Ord)

    data Rigid = RConst Weight
               | RVar RigidId
                 deriving (Eq, Ord, Show)

    type NodeId  = Int
    type RigidId = Int
    type FlexId  = Int
    type Scope   = RigidId -> Bool
    -- which rigid variables a flex may be instatiated to

    instance Show Node where
      show (Flex  i) = "?" ++ show i
      show (Rigid (RVar i)) = "v" ++ show i
      show (Rigid (RConst Infinite))   = "#"
      show (Rigid (RConst (Finite n))) = show n

    infinite (RConst Infinite) = True
    infinite _ = False

    -- isBelow r w r'
    -- checks, if r and r' are connected by w (meaning w not infinite)
    -- wether r + w <= r'
    -- precondition: not the same rigid variable
    isBelow :: Rigid -> Weight -> Rigid -> Bool
    isBelow _ Infinite _ = True
    isBelow _ n (RConst Infinite) = True
    -- isBelow (RConst Infinite)   n (RConst (Finite _)) = False
    isBelow (RConst (Finite i)) (Finite n) (RConst (Finite j)) = i + n <= j
    isBelow _ _ _ = False -- rigid variables are not related

    -- a constraint is an edge in the graph
    data Constraint = NewFlex FlexId Scope
                    | Arc Node Int Node
    -- Arc v1 k v2  at least one of v1,v2 is a VMeta (Flex),
    --              the other a VMeta or a VGen (Rigid)
    -- if k <= 0 this means  $^(-k) v1 <= v2
    -- otherwise                    v1 <= $^k v3

    instance Show Constraint where
      show (NewFlex i s) = "SizeMeta(?" ++ show i ++ ")"
      show (Arc v1 k v2) | k == 0 = show v1 ++ "<=" ++ show v2
                         | k < 0  = show v1 ++ "+" ++ show (-k) ++ "<=" ++ show v2
                         | otherwise  = show v1 ++ "<=" ++ show v2 ++ "+" ++ show k

    type Constraints = [Constraint]

    emptyConstraints = []

    -- graph (matrix) ------------------------------------------------

    data Graph = Graph
      { flexScope :: Map FlexId Scope        -- scope for each flexible var
      , nodeMap :: Map Node NodeId           -- node labels to node numbers
      , intMap  :: Map NodeId Node           -- node numbers to node labels
      , nextNode :: NodeId                   -- number of nodes (n)
      , graph :: NodeId -> NodeId -> Weight  -- the edges (restrict to [0..n[)
      }

    -- the empty graph: no nodes, edges are all undefined (infinity weight)
    initGraph = Graph Map.empty Map.empty Map.empty 0 (\ x y -> Infinite)

    -- the Graph Monad, for constructing a graph iteratively
    type GM = State Graph

    addFlex :: FlexId -> Scope -> GM ()
    addFlex x scope = do
      st <- get
      put $ st { flexScope = Map.insert x scope (flexScope st) }
      addNode (Flex x)
      return ()

    -- i <- addNode n  returns number of node n. if not present, it is added first
    addNode :: Node -> GM Int
    addNode n = do
      st <- get
      case Map.lookup n (nodeMap st) of
        Just i -> return i
        Nothing -> do let i = nextNode st
                      put $ st { nodeMap = Map.insert n i (nodeMap st)
                               , intMap = Map.insert i n (intMap st)
                               , nextNode = i + 1
                               }
                      return i

    -- addEdge n1 k n2
    -- improves the weight of egde n1->n2 to be at most k
    -- also adds nodes if not yet present
    addEdge :: Node -> Int -> Node -> GM ()
    addEdge n1 k n2 = do
      i1 <- addNode n1
      i2 <- addNode n2
      st <- get
      let graph' x y = if (x,y) == (i1,i2) then Finite k `oplus` (graph st) x y
                       else graph st x y
      put $ st { graph = graph' }

    addConstraint :: Constraint -> GM ()
    addConstraint (NewFlex x scope) = addFlex x scope
    addConstraint (Arc n1 k n2)     = addEdge n1 k n2

    buildGraph :: Constraints -> Graph
    buildGraph cs = execState (mapM_ addConstraint cs) initGraph

    mkMatrix :: Int -> (Int -> Int -> Weight) -> Matrix Weight
    mkMatrix n g = array ((0,0),(n-1,n-1))
                     [ ((i,j), g i j) | i <- [0..n-1], j <- [0..n-1]]

    -- displaying matrices with row and column labels --------------------

    -- a matrix with row descriptions in b and column descriptions in c
    data LegendMatrix a b c = LegendMatrix
      { matrix   :: Matrix a
      , rowdescr :: Int -> b
      , coldescr :: Int -> c
      }

    instance (Show a, Show b, Show c) => Show (LegendMatrix a b c) where
      show (LegendMatrix m rd cd) =
        -- first show column description
        let ((r,c),(r',c')) = bounds m
        in foldr (\ j s -> "\t" ++ show (cd j) ++ s) "" [c .. c'] ++
        -- then output rows
           foldr (\ i s -> "\n" ++ show (rd i) ++
                    foldr (\ j t -> "\t" ++ show (m!(i,j)) ++ t)
                          (s)
                          [c .. c'])
                 "" [r .. r']

    -- solving the constraints -------------------------------------------

    -- a solution assigns to each flexible variable a size expression
    -- which is either a constant or a v + n for a rigid variable v
    type Solution = Map Int SizeExpr

    emptySolution = Map.empty
    extendSolution subst k v = Map.insert k v subst

    data SizeExpr = SizeVar Int Int   -- e.g. x + 5
                  | SizeConst Weight  -- a number or infinity

    instance Show SizeExpr where
      show (SizeVar n 0) = show (Rigid (RVar n))
      show (SizeVar n k) = show (Rigid (RVar n)) ++ "+" ++ show k
      show (SizeConst w) = show w

    -- sizeRigid r n  returns the size expression corresponding to r + n
    sizeRigid :: Rigid -> Int -> SizeExpr
    sizeRigid (RConst k) n = SizeConst (inc k n)
    sizeRigid (RVar i)   n = SizeVar i n

    {-
    apply :: SizeExpr -> Solution -> SizeExpr
    apply e@(SizeExpr (Rigid _) _) phi = e
    apply e@(SizeExpr (Flex  x) i) phi = case Map.lookup x phi of
      Nothing -> e
      Just (SizeExpr v j) -> SizeExpr v (i + j)

    after :: Solution -> Solution -> Solution
    after psi phi = Map.map (\ e -> e `apply` phi) psi
    -}

    {-
    solve :: Constraints -> Maybe Solution
    solve cs = if any (\ x -> x < Finite 0) d then Nothing
         else Map.
       where gr = buildGraph cs
             n  = nextNode gr
             m  = mkMatrix n (graph gr)
             m' = warshall m
             d  = [ m!(i,i) | i <- [0 .. (n-1)] ]
             ns = keys (nodeMap gr)
    -}

    {- compute solution

    a solution CANNOT exist if

      v < v  for a rigid variable v

      v <= v' for rigid variables v,v'

      x < v   for a flexible variable x and a rigid variable v

    thus, for each flexible x, only one of the following cases is possible

      r+n <= x+m <= infty  for a unique rigid r  (meaning r --(m-n)--> x)
      x <= r+n             for a unique rigid r  (meaning x --(n)--> r)

    we are looking for the least values for flexible variables that solve
    the constraints.  Algorithm

    while flexible variables and rigid rows left
      find a rigid variable row i
        for all flexible columns j
          if i --n--> j with n<=0 (meaning i+n <= j) then j = i + n

    while flexible variables j left
      search the row j for entry i
        if j --n--> i with n >= 0 (meaning j <= i + n) then j = i + n

    -}

    solve :: Constraints -> Maybe Solution
    solve cs = -- trace (show cs) $
       -- trace (show lm0) $
        -- trace (show lm) $ -- trace (show d) $
         let solution = if solvable then loop1 flexs rigids emptySolution
                        else Nothing
         in -- trace (show solution) $
             solution
       where -- compute the graph and its transitive closure m
             gr  = buildGraph cs
             n   = nextNode gr            -- number of nodes
             m0  = mkMatrix n (graph gr)
             m   = warshall m0

             -- tracing only: build output version of transitive graph
             legend i = fromJust $ Map.lookup i (intMap gr) -- trace only
             lm0 = LegendMatrix m0 legend legend            -- trace only
             lm  = LegendMatrix m legend legend             -- trace only

             -- compute the sets of flexible and rigid node numbers
             ns  = Map.keys (nodeMap gr)
             -- a set of flexible variables
             flexs  = foldl (\ l k -> case k of (Flex i) -> i : l
                                                (Rigid _) -> l) [] ns
             -- a set of rigid variables
             rigids = foldl (\ l k -> case k of (Flex _) -> l
                                                (Rigid i) -> i : l) [] ns

             -- rigid matrix indices
             rInds = foldl (\ l r -> let Just i = Map.lookup (Rigid r) (nodeMap gr)
                                     in i : l) [] rigids

             -- check whether there is a solution
             -- d   = [ m!(i,i) | i <- [0 .. (n-1)] ]  -- diagonal
    -- a rigid variable might not be less than it self, so no -.. on the
    -- rigid part of the diagonal
             solvable = all (\ x -> x >= Finite 0) [ m!(i,i) | i <- rInds ] &&
    -- a rigid variable might not be bounded below by infinity or
    -- bounded above by a constant
    -- it might not be related to another rigid variable
               all (\ (r,  r') -> r == r' ||
                    let Just row = (Map.lookup (Rigid r)  (nodeMap gr))
                        Just col = (Map.lookup (Rigid r') (nodeMap gr))
                        edge = m!(row,col)
                    in  isBelow r edge r' )
                 [ (r,r') | r <- rigids, r' <- rigids ]
               &&
    -- a flexible variable might not be strictly below a rigid variable
               all (\ (x, v) ->
                    let Just row = (Map.lookup (Flex x)  (nodeMap gr))
                        Just col = (Map.lookup (Rigid (RVar v)) (nodeMap gr))
                        edge = m!(row,col)
                    in  edge >= Finite 0)
                 [ (x,v) | x <- flexs, (RVar v) <- rigids ]

             inScope :: FlexId -> Rigid -> Bool
             inScope x (RConst _) = True
             inScope x (RVar v)   = scope v
                 where Just scope = Map.lookup x (flexScope gr)

    {- loop1

    while flexible variables and rigid rows left
      find a rigid variable row i
        for all flexible columns j
          if i --n--> j with n<=0 (meaning i + n <= j) then j = i + n

    -}
             loop1 :: [FlexId] -> [Rigid] -> Solution -> Maybe Solution
             loop1 [] rgds subst = Just subst
             loop1 flxs [] subst = loop2 flxs subst
             loop1 flxs (r:rgds) subst =
                let row = fromJust $ Map.lookup (Rigid r) (nodeMap gr)
                    (flxs',subst') =
                      foldl (\ (flx,sub) f ->
                              let col = fromJust $ Map.lookup (Flex f) (nodeMap gr)
                              in  case (inScope f r, m!(row,col)) of
    --                                Finite z | z <= 0 ->
                                    (True, Finite z) ->
                                       let trunc z | z >= 0 = 0
                                                | otherwise = -z
                                       in (flx, extendSolution sub f (sizeRigid r (trunc z)))
                                    _ -> (f : flx, sub)
                         ) ([], subst) flxs
                in loop1 flxs' rgds subst'

    {- loop2

    while flexible variables j left
      search the row j for entry i
        if j --n--> i with n >= 0 (meaning j <= i + n) then j = i

    -}
             loop2 :: [FlexId] -> Solution -> Maybe Solution
             loop2 [] subst = Just subst
             loop2 (f:flxs) subst = loop3 0 subst
               where row = fromJust $ Map.lookup (Flex f) (nodeMap gr)
                     loop3 col subst | col >= n =
                       -- default to infinity
                        loop2 flxs (extendSolution subst f (SizeConst Infinite))
                     loop3 col subst =
                       case Map.lookup col (intMap gr) of
                         Just (Rigid r) | not (infinite r) ->
                           case (inScope f r, m!(row,col)) of
                            (True, Finite z) | z >= 0 ->
                                loop2 flxs (extendSolution subst f (sizeRigid r z))
                            (_, Infinite) -> loop3 (col+1) subst
                            _ -> -- trace ("unusable rigid: " ++ show r ++ " for flex " ++ show f)
                                  Nothing  -- NOT: loop3 (col+1) subst
                         _ -> loop3 (col+1) subst

    -- Testing ----------------------------------------------------------------

    genGraph :: Ord node => Float -> Gen edge -> [node] -> Gen (AdjList node edge)
    genGraph density edge nodes = do
      Map.fromList . concat <$> mapM neighbours nodes
      where
        k = round (100 * density)
        neighbours n = do
          ns <- concat <$> mapM neighbour nodes
          case ns of
            []  -> elements [[(n, [])], []]
            _   -> return [(n, ns)]
        neighbour n = frequency
          [ (k, do e <- edge
                   ns <- neighbour n
                   return ((n, e):ns))
          , (100 - k, return [])
          ]

    newtype Distance = Dist Nat
      deriving (Eq, Ord, Num, Integral, Show, Enum, Real)

    instance SemiRing Distance where
      oplus  (Dist a) (Dist b) = Dist (min a b)
      otimes (Dist a) (Dist b) = Dist (a + b)

    genGraph_ :: Nat -> Gen (AdjList Nat Distance)
    genGraph_ n =
      genGraph 0.2 (Dist <$> natural) [0..n - 1]

    lookupEdge :: Ord n => n -> n -> AdjList n e -> Maybe e
    lookupEdge i j g = lookup j =<< Map.lookup i g

    edges :: Ord n => AdjList n e -> [(n,n,e)]
    edges g = do
      (i, ns) <- Map.toList g
      (j, e)  <- ns
      return (i, j, e)

    -- | Check that no edges get longer when completing a graph.
    prop_smaller n' =
      forAll (genGraph_ n) $ \g ->
      let g' = warshallG g in
      and [ lookupEdge i j g' =< e
          | (i, j, e) <- edges g
          ]
      where
        n = abs (div n' 2)
        Nothing =< _ = False
        Just x  =< y = x <= y

    newEdge i j e = Map.insertWith (++) i [(j, e)]

    genPath :: Nat -> Nat -> Nat -> AdjList Nat Distance -> Gen (AdjList Nat Distance)
    genPath n i j g = do
      es <- listOf $ (,) <$> node <*> edge
      v  <- edge
      return $ addPath i (es ++ [(j, v)]) g
      where
        edge = Dist <$> natural
        node = choose (0, n - 1)
        addPath _ [] g = g
        addPath i ((j, v):es) g =
          newEdge i j v $ addPath j es g

    -- | Check that all transitive edges are added.
    prop_path n' =
      forAll (genGraph_ n) $ \g ->
      forAll (replicateM 2 $ choose (0, n - 1)) $ \[i, j] ->
      forAll (genPath n i j g) $ \g' ->
      isJust (lookupEdge i j $ warshallG g')
      where
        n = abs (div n' 2) + 1

    mapNodes :: (Ord node, Ord node') => (node -> node') -> AdjList node edge -> AdjList node' edge
    mapNodes f = Map.map f' . Map.mapKeys f
      where
        f' es = [ (f n, e) | (n,e) <- es ]

    -- | Check that no edges are added between components.
    prop_disjoint n' =
      forAll (replicateM 2 $ genGraph_ n) $ \[g1, g2] ->
      let g  = Map.union (mapNodes Left g1) (mapNodes Right g2)
          g' = warshallG g
      in all disjoint (Map.assocs g')
      where
        n = abs (div n' 3)
        disjoint (Left i, es)  = all (isLeft . fst) es
        disjoint (Right i, es) = all (isRight . fst) es
        isLeft = either (const True) (const False)
        isRight = not . isLeft

    prop_stable n' =
      forAll (genGraph_ n) $ \g ->
      let g' = warshallG g in
      g' =~= warshallG g'
      where
        n = abs (div n' 2)
        g =~= g' = sort (edges g) == sort (edges g')

    tests :: IO Bool
    tests = runTests "Agda.Utils.Warshall"
      [ quickCheck' prop_smaller
      , quickCheck' prop_path
      , quickCheck' prop_disjoint
      , quickCheck' prop_stable
      ]