module Johnson(johnson_all_pairs_shortest_paths) where
  import Graph
  import PropertyMap
  import BellmanFord
  import Dijkstra
  import AdjacencyList
  
  johnson_all_pairs_shortest_paths ::
    (VertexListGraph g v, EdgeListGraph g e v, Num n, Ord n, Ord v,
     ReadPropertyMap i v Int,  ReadPropertyMap w e n) =>
    g -> w -> i -> n -> n -> Maybe (FMap v (FMap v n))
  johnson_all_pairs_shortest_paths g1 w0 vidx zero inf =
    let n = (num_vertices g1)
        idx2vert = create_map (zip [0..(n-1)] (vertices g1)) (head (vertices g1))
        g2 = adj_list (1 + n)
          ([((get vidx (src e g1)), (get vidx (tgt e g1))) | e <- edges g1]
            ++ [(n, (get vidx v)) | v <- vertices g1])
        w1 = (create_map ([((edge (get vidx (src e g1)) (get vidx (tgt e g1))), (get w0 e)) | e <- edges g1 ]
                          ++ [(edge n (get vidx v), zero) | v <- vertices g1]) zero)
        s = vertex n
        d1 = (put (init_map (vertices g2) inf) s zero)
        p1 = (create_map [(v,v) | v <- vertices g2 ] s)
        (d2, p2, neg_cycle) 
          = (bellman_ford_shortest_paths g2 (n+1) w1 p1 d1 (+) (<))
    in if (neg_cycle) then Nothing
       else
         let h = (create_map [(v, (get d2 v)) | v <- vertices g2 ] zero)
             w_hat = (create_map
 		       [(e, (get w1 e) + (get h (src e g2))
                           - (get h (tgt e g2)))
		        | e <- edges g2 ] zero)
         in Just (create_map
              [ (let (d3,p3) = (dijkstra_shortest_paths g2 u d2 w_hat (<) (+) zero)
                 in (get idx2vert (idx u), 
                     (create_map [((get idx2vert (idx v)), (get d3 v))
                                 | v <- vertices g2, (idx v) < n ] zero)))
               | u <- vertices g2 ] (init_map (vertices g1) zero))